jax-ml/jax · error · ConcretizationTypeError

The error occurred in the __reduce__ method, which may indic

Error message

The error occurred in the __reduce__ method, which may indicate an attempt to serialize/pickle a traced value.

What it means

Raised by JAXTracer.__reduce__ when something attempts to pickle or deep-copy a JAX Tracer. Tracers hold references to live trace state (the transformation currently recording them) and have no concrete value, so they cannot be serialized. JAX deliberately raises ConcretizationTypeError with this message to point at accidental serialization of traced values.

Source

Thrown at jax/_src/core.py:1155

    return self.aval._hex(self)

  def __oct__(self):
    if is_concrete(self): return oct(self.to_concrete_value())
    check_integer_conversion(self)
    if not hasattr(self.aval, "_oct"):
      raise TypeError(f"Value of type {type(self)} is not convertible to oct.")
    return self.aval._oct(self)

  def __index__(self):
    if is_concrete(self): return operator.index(self.to_concrete_value())
    check_integer_conversion(self)
    if not hasattr(self.aval, "_index"):
      raise TypeError(f"Value of type {type(self)} is not convertible to integer index.")
    return self.aval._index(self)

  # raises a useful error on attempts to pickle a Tracer.
  def __reduce__(self):
    raise ConcretizationTypeError(
      self, ("The error occurred in the __reduce__ method, which may "
             "indicate an attempt to serialize/pickle a traced value."))

  # raises the better error message from ShapedArray
  def __setitem__(self, key, value):
    if not hasattr(self.aval, "_setitem"):
      raise TypeError(f"Value of type {type(self)} is not indexable.")
    return self.aval._setitem(self, key, value)

  # NumPy also only looks up special methods on classes.
  def __array_module__(self, types):
    if not hasattr(self.aval, "_array_module"):
      raise TypeError(f"Value of type {type(self)} is not compatible with the Array API.")
    return self.aval._array_module(self, types)

  def __getattr__(self, name):
    # if the aval property raises an AttributeError, gets caught here
    assert not config.enable_checks.value or name != "aval"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure only concrete jax.Array (or NumPy/Python) values are pickled: convert with np.asarray(x) / jax.device_get(x) before serialization.
2. Remove or restructure caching (lru_cache/joblib.Memory) so it does not key or store on objects created inside jit/grad.
3. For multiprocessing, pass concrete arrays and re-jit in the worker; or switch to a process pool that supports shared memory for JAX arrays.
4. If deep-copying, copy the underlying concrete data, not objects holding tracers.

Example fix

# before
@jax.jit
def f(x):
    obj = {'params': x}          # x is a Tracer
    pickle.dumps(obj)            # ConcretizationTypeError in __reduce__

# after
@jax.jit
def f(x):
    return x * 2                 # return tracer
out = f(x)
pickle.dumps(np.asarray(out))   # pickle concrete value on host

import jax, numpy as np
def pickle_safe(obj):
    def concrete(x):
        return np.asarray(jax.device_get(x)) if isinstance(x, (jax.Array, jax.core.Tracer)) else x
    # recursively concretize before dumping
    return pickle.dumps(concrete(obj))

import jax
def contains_tracer(obj) -> bool:
    return any(isinstance(v, jax.core.Tracer) for v in jax.tree_util.tree_leaves(obj))

try:
    payload = pickle.dumps(obj)
except jax.errors.ConcretizationTypeError:
    obj = jax.tree_util.tree_map(lambda x: np.asarray(x), obj)
    payload = pickle.dumps(obj)

When it happens

Trigger: Calling pickle.dumps / cloudpickle / copy.deepcopy on an object graph that contains a Tracer; caching libraries (functools.lru_cache wrappers around traced calls, joblib, diskcache) that pickle arguments; multiprocessing that pickles captured state; returning a Tracer and storing it in a serialized cache or queue.

Common situations: functools.cache on a function that receives jitted/traced arrays; sending JAX arrays or closures holding them through multiprocessing or ray workers that pickle payloads; checkpointing frameworks that pickle model state that accidentally includes a tracer from inside a grad pass; deep-copying configs that captured traced values.

Related errors

• The names should be exclusive and should not intersect in `n
• stop_gradient only works on valid JAX arrays, but input argu
• {self.__class__.__name__} has no attribute {name}
• Value of type {type(self)} is not convertible to float.
• Value of type {type(self)} is not convertible to complex.

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