jax-ml/jax · error · ConcretizationTypeError

The is_fully_replicated property was called on {self._error_

Error message

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

What it means

ConcretizationTypeError raised when the Tracer.is_fully_replicated property is accessed. is_fully_replicated reports whether all shards of a materialized jax.Array hold identical data (relevant to sharded/distributed arrays); tracers carry no layout, so the property stub raises during tracing.

Source

Thrown at jax/_src/core.py:1267

  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()}."
      f"{self._origin_msg()}")

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the replication check outside the traced function and pass the result as a static boolean.
Derive replication info from the value's sharding on its aval (e.g. via jax.typeof) where tracing-time knowledge is required.
Let the XLA/SPMD compiler handle replication-aware elision of collectives rather than branching in Python.
Guard with isinstance(x, jax.Array) in shared helper code.

Example fix

// before
@jax.jit
def allreduce(x):
    if x.is_fully_replicated:  # ConcretizationTypeError
        return x
    return lax.psum(x, 'i')

// after
def allreduce(x):
    if x.is_fully_replicated:   # concrete check outside
        return x
    return _psum_jit(x)

Defensive patterns

Strategy: validation

Validate before calling

import jax

def fully_replicated_or_default(x, default=False):
    if isinstance(x, jax.core.Tracer):
        return default
    return x.is_fully_replicated

Type guard

import jax
from jax.core import Tracer

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

Try / catch

from jax.errors import ConcretizationTypeError
try:
    rep = x.is_fully_replicated
except ConcretizationTypeError:
    rep = False  # conservative fallback under tracing

Prevention

Hoist replication checks to the caller; pass as static boolean.
Let the compiler elide redundant collectives instead of Python branching.
Document which helper functions are trace-safe vs host-only.

When it happens

Trigger: Accessing `x.is_fully_replicated` on a traced value inside jit/pjit/grad/vmap/scan/remat or inside a custom transform rule.

Common situations: Optimization branches like `if x.is_fully_replicated: skip_all_reduce(x)` placed inside a jitted collective function; distributed training code refactored under vmap; sharding assertions copied into traced kernels.

Related errors

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