{"record":{"id":"59aa23e11bd92bee","repo":"jax-ml/jax","slug":"the-error-occurred-in-the-reduce-method-which","errorCode":null,"errorMessage":"The error occurred in the __reduce__ method, which may indicate an attempt to serialize/pickle a traced value.","messagePattern":"The error occurred in the __reduce__ method, which may indicate an attempt to serialize/pickle a traced value\\.","errorType":"exception","errorClass":"ConcretizationTypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":1155,"sourceCode":"    return self.aval._hex(self)\n\n  def __oct__(self):\n    if is_concrete(self): return oct(self.to_concrete_value())\n    check_integer_conversion(self)\n    if not hasattr(self.aval, \"_oct\"):\n      raise TypeError(f\"Value of type {type(self)} is not convertible to oct.\")\n    return self.aval._oct(self)\n\n  def __index__(self):\n    if is_concrete(self): return operator.index(self.to_concrete_value())\n    check_integer_conversion(self)\n    if not hasattr(self.aval, \"_index\"):\n      raise TypeError(f\"Value of type {type(self)} is not convertible to integer index.\")\n    return self.aval._index(self)\n\n  # raises a useful error on attempts to pickle a Tracer.\n  def __reduce__(self):\n    raise ConcretizationTypeError(\n      self, (\"The error occurred in the __reduce__ method, which may \"\n             \"indicate an attempt to serialize/pickle a traced value.\"))\n\n  # raises the better error message from ShapedArray\n  def __setitem__(self, key, value):\n    if not hasattr(self.aval, \"_setitem\"):\n      raise TypeError(f\"Value of type {type(self)} is not indexable.\")\n    return self.aval._setitem(self, key, value)\n\n  # NumPy also only looks up special methods on classes.\n  def __array_module__(self, types):\n    if not hasattr(self.aval, \"_array_module\"):\n      raise TypeError(f\"Value of type {type(self)} is not compatible with the Array API.\")\n    return self.aval._array_module(self, types)\n\n  def __getattr__(self, name):\n    # if the aval property raises an AttributeError, gets caught here\n    assert not config.enable_checks.value or name != \"aval\"","sourceCodeStart":1137,"sourceCodeEnd":1173,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L1137-L1173","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Ensure only concrete jax.Array (or NumPy/Python) values are pickled: convert with np.asarray(x) / jax.device_get(x) before serialization.","Remove or restructure caching (lru_cache/joblib.Memory) so it does not key or store on objects created inside jit/grad.","For multiprocessing, pass concrete arrays and re-jit in the worker; or switch to a process pool that supports shared memory for JAX arrays.","If deep-copying, copy the underlying concrete data, not objects holding tracers."],"exampleFix":"# before\n@jax.jit\ndef f(x):\n    obj = {'params': x}          # x is a Tracer\n    pickle.dumps(obj)            # ConcretizationTypeError in __reduce__\n\n# after\n@jax.jit\ndef f(x):\n    return x * 2                 # return tracer\nout = f(x)\npickle.dumps(np.asarray(out))   # pickle concrete value on host","handlingStrategy":"validation","validationCode":"import jax, numpy as np\ndef pickle_safe(obj):\n    def concrete(x):\n        return np.asarray(jax.device_get(x)) if isinstance(x, (jax.Array, jax.core.Tracer)) else x\n    # recursively concretize before dumping\n    return pickle.dumps(concrete(obj))","typeGuard":"import jax\ndef contains_tracer(obj) -> bool:\n    return any(isinstance(v, jax.core.Tracer) for v in jax.tree_util.tree_leaves(obj))","tryCatchPattern":"try:\n    payload = pickle.dumps(obj)\nexcept jax.errors.ConcretizationTypeError:\n    obj = jax.tree_util.tree_map(lambda x: np.asarray(x), obj)\n    payload = pickle.dumps(obj)","preventionTips":["Never cache (lru_cache/joblib) results captured inside jit/grad; cache concrete outputs only.","Concretize with jax.device_get / np.asarray at transform boundaries before storing or sending values.","In multiprocessing, ship concrete arrays and re-trace in the worker.","Deep-copy plain data, not objects that may hold tracers."],"tags":["jax","tracer","pickle","serialization","multiprocessing","caching"],"backgroundTag":"object-not-picklable","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}