jax-ml/jax · error · ConcretizationTypeError

The on_device_size_in_bytes() method was called on {self._er

Error message

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

What it means

ConcretizationTypeError raised when Tracer.on_device_size_in_bytes() is called. This method reports the device-side byte size of a materialized jax.Array's buffers; a tracer has no storage, so the base-class stub raises during any JAX transform's tracing.

Source

Thrown at jax/_src/core.py:1272

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

_jax.set_tracer_class(Tracer)

# these can be used to set up forwarding of properties and instance methods from
# Tracer instances to the underlying avals

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Compute size accounting outside the trace, on the concrete arrays passed to/returned from jit.
Estimate sizes statically from shapes/dtypes (x.aval.shape, itemsize) if you need them during tracing.
Remove instrumentation from traced kernels and collect it in the surrounding Python driver.
Cache per-shape byte sizes keyed by input shapes to avoid repeated host queries.

Example fix

// before
@jax.jit
def step(x):
    log(x.on_device_size_in_bytes())  # ConcretizationTypeError
    return heavy(x)

// after
def step_outer(x):
    log(x.on_device_size_in_bytes())  # concrete array
    return _step_jit(x)

Defensive patterns

Strategy: type-guard

Validate before calling

import jax

def device_bytes(x):
    if isinstance(x, jax.core.Tracer):
        import numpy as np
        a = x.aval
        return int(np.prod(a.shape)) * a.dtype.itemsize  # static estimate
    return x.on_device_size_in_bytes()

Type guard

import jax
from jax.core import Tracer

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

Try / catch

from jax.errors import ConcretizationTypeError
try:
    nbytes = x.on_device_size_in_bytes()
except ConcretizationTypeError:
    nbytes = None  # use static shape-based estimate instead

Prevention

Do memory accounting in the host driver on concrete arrays.
Estimate sizes from shape/dtype when inside transforms.
Keep telemetry calls out of jitted kernels.

When it happens

Trigger: Calling `x.on_device_size_in_bytes()` on a value traced by jit/grad/vmap/pmap/scan/remat or a custom rule.

Common situations: Memory-budget accounting or telemetry (`total += x.on_device_size_in_bytes()`) wired into a jitted step; profiling helpers reused under vmap; OOM-debugging instrumentation left in traced code.

Related errors

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