jax-ml/jax · error · ConcretizationTypeError

The is_fully_addressable property was called on {self._error

Error message

The is_fully_addressable property was called on {self._error_repr()}.{self._origin_msg()}

What it means

ConcretizationTypeError raised when the Tracer.is_fully_addressable property is accessed. is_fully_addressable indicates whether every shard of a materialized (potentially SPMD) jax.Array is addressable by the local process; tracers have no distributed layout, so the property stub raises during tracing.

Source

Thrown at jax/_src/core.py:1261

  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()}")

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

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

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

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Hoist the is_fully_addressable check outside the trace and pass the boolean as a static argument to the jitted function.
Branch on jax.typeof(x)/sharding info available on the aval if you must decide inside tracing.
Restructure so distributed-awareness lives in the outer (non-traced) orchestration code, not in the kernel.
Use functools.partial or closure over the precomputed flag.

Example fix

// before
@jax.jit
def f(x):
    if x.is_fully_addressable:   # ConcretizationTypeError
        return local_path(x)
    return dist_path(x)

// after
def f(x):
    addr = x.is_fully_addressable          # concrete array
    return _f_jit(x, addr)
@partial(jax.jit, static_argnames=['addr'])
def _f_jit(x, addr):
    return local_path(x) if addr else dist_path(x)

Defensive patterns

Strategy: validation

Validate before calling

import jax

def fully_addressable_or_default(x, default=True):
    if isinstance(x, jax.core.Tracer):
        return default
    return x.is_fully_addressable

Type guard

import jax
from jax.core import Tracer

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

Try / catch

from jax.errors import ConcretizationTypeError
try:
    addr = x.is_fully_addressable
except ConcretizationTypeError:
    addr = None  # decide via a static argument instead

Prevention

Compute addressability once on the host and pass it via static_argnums.
Keep distributed branching out of traced kernels.
Write integration tests that run the same code under jit and eagerly.

When it happens

Trigger: Reading `x.is_fully_addressable` on a traced value inside jit/pjit/grad/vmap/scan or a custom rule; commonly via an assert or conditional on it.

Common situations: Single-host vs multi-host branching logic (`if x.is_fully_addressable:`) written for distributed arrays and then executed inside jitted code; tests running under vmap that hit distributed-guard code; refactoring multi-process pipeline code into transformed functions.

Related errors

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