jax-ml/jax · error · AttributeError

The 'sharding' attribute is not available on {self._error_re

Error message

The 'sharding' attribute is not available on {self._error_repr()}. To query sharding information on tracers, use `jax.typeof(x)`.

What it means

Raised by JAXTracer.__getattr__ when code accesses the .sharding attribute on a JAX Tracer. Sharding describes how a concrete, dispatched array's data is laid out across devices; during tracing (jit/pmap/vmap/shard_map) the value is abstract and has no committed sharding, so the attribute does not exist and lookup raises AttributeError with a pointer to jax.typeof.

Source

Thrown at jax/_src/core.py:1183

  # NumPy also only looks up special methods on classes.
  def __array_module__(self, types):
    if not hasattr(self.aval, "_array_module"):
      raise TypeError(f"Value of type {type(self)} is not compatible with the Array API.")
    return self.aval._array_module(self, types)

  def __getattr__(self, name):
    # if the aval property raises an AttributeError, gets caught here
    assert not config.enable_checks.value or name != "aval"

    # These must raise AttributeError in the base class for backward compatibility.
    # TODO(jakevdp): can we change this and make them raise NotImplementedError instead?
    if name in ["block_until_ready", "copy_to_host_async"]:
      raise AttributeError(
        f"The '{name}' method is not available on {self._error_repr()}."
        f"{self._origin_msg()}")

    if name == 'sharding':
      raise AttributeError(
        f"The 'sharding' attribute is not available on {self._error_repr()}. "
        "To query sharding information on tracers, use `jax.typeof(x)`.")

    try:
      attr = getattr(self.aval, name)
    except AttributeError as err:
      raise AttributeError(
          f"{self.__class__.__name__} has no attribute {name}"
      ) from err
    else:
      t = type(attr)
      if t is aval_property:
        return attr.fget(self)
      elif t is aval_method:
        return types.MethodType(attr.fun, self)
      else:
        return attr

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Query sharding outside the traced function on concrete arrays, or use jax.typeof(x) inside traces as the message suggests to get abstract type info.
  2. Pass sharding knowledge explicitly (e.g. via in_shardings/out_shardings on jax.jit or jax.lax.with_sharding_constraint) instead of introspecting .sharding.
  3. Guard generic introspection with isinstance(x, jax.Array) before touching .sharding.
  4. For pmap-era code, migrate to jax.jit with sharding-in annotations.

Example fix

# before
@jax.jit
def step(x):
    nd = x.sharding.num_devices   # AttributeError on Tracer
    return x / nd

# after
@jax.jit
def step(x):
    t = jax.typeof(x)             # abstract type available on tracers
    return x / t.size  # or hoist num_devices out as a static scalar
Defensive patterns

Strategy: type-guard

Validate before calling

import jax
def sharding_or_none(x):
    if isinstance(x, jax.core.Tracer):
        return None  # no sharding during tracing; use jax.typeof(x) if type info needed
    return x.sharding

Type guard

import jax
def has_concrete_sharding(x) -> bool:
    return isinstance(x, jax.Array) and not isinstance(x, jax.core.Tracer)

Try / catch

try:
    s = x.sharding
except AttributeError:
    s = None  # tracer inside jit/pmap; use jax.typeof(x) for abstract type

Prevention

When it happens

Trigger: Reading x.sharding inside a jitted/pmap/shard_map-decorated function; generic code (e.g. MultiDeviceHost or data-parallel wrappers) that inspects .sharding on every array-like argument; checks like x.sharding.num_devices or NamedSharding queries applied to traced values.

Common situations: Data-parallel training loops that auto-shard based on incoming arrays' sharding, invoked on tracers; migration from pmap (where .sharding probes sometimes slipped through) to jax.jit + sharding constraints; logging or assertion helpers inspecting device placement of any array they touch.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/5b36306151f358b7. Report an issue: GitHub.