jax-ml/jax · error · ConcretizationTypeError

The addressable_data() method was called on {self._error_rep

Error message

The addressable_data() method was called on {self._error_repr()}.{self._origin_msg()}

What it means

ConcretizationTypeError raised when Tracer.addressable_data(index) is called. addressable_data is only implemented for materialized arrays (it returns the per-process shard of a sharded jax.Array); tracers carry no buffers, so during tracing (jit/grad/vmap/etc.) the base-class stub raises.

Source

Thrown at jax/_src/core.py:1234

          for name, pp_payload in contents])
      ])))
    return base

  def __repr__(self):
    return self._pretty_print(verbose=False).format()

  def _contents(self):
    try:
      return [(name, getattr(self, name)) for name in self.__slots__]
    except AttributeError:
      return ()

  def _origin_msg(self) -> str:
    return ""

  # Methods that are only valid for materialized arrays
  def addressable_data(self, index):
    raise ConcretizationTypeError(self,
      f"The addressable_data() method was called on {self._error_repr()}."
      f"{self._origin_msg()}")

  def delete(self):
    raise ConcretizationTypeError(self,
      f"The delete() method was called on {self._error_repr()}."
      f"{self._origin_msg()}")

  def devices(self):
    raise ConcretizationTypeError(self,
      f"The devices() method was called on {self._error_repr()}."
      f"{self._origin_msg()}")

  @property
  def global_shards(self):
    raise ConcretizationTypeError(self,
      f"The global_shards property was called on {self._error_repr()}."
      f"{self._origin_msg()}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the addressable_data() call outside the traced function: compute shard layout before entering jit and pass results as static/closed-over constants.
If you need per-shard data under SPMD, restructure to use sharding-aware ops (e.g. pjit with in_shardings, jax.lax ops, or addressable data computed after the jit call returns).
Guard with isinstance(x, jax.Array) before calling, and take a tracing-time fallback path.
If you intended a concrete value but got a tracer, check for an accidentally nested jit/grad wrapping (e.g. grad-of-jit where the inner call is traced).

Example fix

// before
@jax.jit
def f(x):
    local = x.addressable_data(0)  # ConcretizationTypeError
    return local.sum()

// after
@jax.jit
def f(x):
    return x.sum()
# inspect shards outside tracing:
# y.addressable_data(0)

Defensive patterns

Strategy: type-guard

Validate before calling

import jax

def call_addressable_data(x, index):
    if isinstance(x, jax.core.Tracer):
        raise ValueError('x is traced; inspect shards outside jit/grad/vmap')
    return x.addressable_data(index)

Type guard

from jax.core import Tracer
import jax

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

Try / catch

from jax.errors import ConcretizationTypeError
try:
    shard = x.addressable_data(0)
except ConcretizationTypeError:
    shard = None  # traced value; handle shard logic on the host instead

Prevention

Design shard-inspection as a host-side step around the traced kernel.
Static-typing hints (jax.Array) plus review of jit bodies for array-only methods.
Keep a checklist of array-only APIs (addressable_data, devices, global_shards...) banned inside traces.

When it happens

Trigger: Calling `x.addressable_data(i)` on a value that is a JAX Tracer, i.e. inside jax.jit, jax.grad, vmap, pmap, lax.scan, remat, or custom derivative rules.

Common situations: Code written against sharded pjit/SPMD arrays is reused inside a jitted function; inspection/debug helper that dumps shards is accidentally called under a transform; sharded checkpoint restore code path runs under vmap.

Related errors

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