jax-ml/jax · error · ConcretizationTypeError

The devices() method was called on {self._error_repr()}.{sel

Error message

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

What it means

ConcretizationTypeError raised when Tracer.devices() is called. devices() returns the physical devices a materialized array's buffers live on; a tracer has no buffers and no placement, so during tracing the base-class stub raises ConcretizationTypeError.

Source

Thrown at jax/_src/core.py:1244

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move device checks outside the traced function and perform them on the concrete input/output arrays.
Inside traces, use sharding info from jax.typeof(x) or the aval if placement metadata is needed.
If the check must run per-call, do it on the arguments before calling the jitted function.
Remove the introspection call from hot traced code; it cannot influence XLA compilation anyway.

Example fix

// before
@jax.jit
def f(x):
    print(x.devices())  # ConcretizationTypeError
    return x + 1

// after
def f_outer(x):
    print(x.devices())   # concrete array
    return jitted_add_one(x)

Defensive patterns

Strategy: type-guard

Validate before calling

import jax

def get_devices(x):
    if isinstance(x, jax.core.Tracer):
        return None  # or derive from jax.typeof(x) sharding
    return x.devices()

Type guard

import jax
from jax.core import Tracer

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

Try / catch

from jax.errors import ConcretizationTypeError
try:
    devs = x.devices()
except ConcretizationTypeError:
    devs = None  # placement unknown under tracing

Prevention

Validate placement on concrete inputs before invoking jit.
Use jax.typeof(x) for tracing-time layout info.
Keep device assertions in test/setup code, not in kernels.

When it happens

Trigger: Calling `x.devices()` on a value while it is being traced by jax.jit, jax.grad, vmap, pmap, lax.scan, remat, or a custom transform rule.

Common situations: Device-placement assertions or logging (`assert x.devices()[0].platform == 'gpu'`) copied into jitted code; multi-host SPMD code checking replica placement inside pjit; debug utilities that introspect device residency run under a transform.

Related errors

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