jax-ml/jax · error · ConcretizationTypeError

The unsafe_buffer_pointer() method was called on

Error message

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

What it means

Raised by JAX when unsafe_buffer_pointer() is called on a Tracer (an abstract, not-yet-computed value) during tracing, e.g. inside jit/vmap/grad. The tracer has no concrete memory address because the underlying array data does not exist until the traced computation is compiled and executed.

Solutions

  1. Move the pointer access outside the traced function and pass the concrete array instead
  2. Use static_argnums/static_argnames so the offending argument stays concrete
  3. If the value must be concrete inside the trace, hoist the computation: compute pointer after jit returns the compiled result
  4. Replace buffer-pointer logic with JAX-native APIs (e.g., jax.dlpack for interop)

Example fix

// before
@jax.jit
def f(x):
    return do_something_with_ptr(x.unsafe_buffer_pointer())

// after
@jax.jit
def f(x):
    return jnp.sum(x)
y = f(x)  # concrete now
ptr = y.unsafe_buffer_pointer()
Defensive patterns

Strategy: validation

Validate before calling

import jax
def is_concrete_array(x):
    return not isinstance(x, jax.core.Tracer) and hasattr(x, 'addressable_data')

Prevention

When it happens

Trigger: Calling .unsafe_buffer_pointer() (or passing an array to something that calls it, like some C-extension/CFFI code or custom buffer protocols) on an array that is a Tracer inside jax.jit, vmap, grad, or another transformation.

Common situations: Passing jitted arrays into external libraries expecting raw pointers (e.g., custom CUDA kernels, PyTorch interop, NumPy C-API consumers); using .data_ptr()-style APIs inside functions decorated with @jax.jit.

Related errors


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

Appendix: source

Thrown at jax/_src/core.py:1283

  @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
aval_property = namedtuple("aval_property", ["fget"])
aval_method = namedtuple("aval_method", ["fun"])

pytype_aval_mappings[Tracer] = lambda x: x.aval
dtypes.register_canonicalize_value_handler(Tracer, None)

def check_eval_args(args):
  for arg in args:
    if isinstance(arg, Tracer):
      raise escaped_tracer_error(arg)

View on GitHub (pinned to 1e1c6a8fc0)