{"record":{"id":"15a4732cf29572ca","repo":"jax-ml/jax","slug":"stop-gradient-only-works-on-valid-jax-arrays-but","errorCode":null,"errorMessage":"stop_gradient only works on valid JAX arrays, but input argument is: {x}","messagePattern":"stop_gradient only works on valid JAX arrays, but input argument is: (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/ad_util.py","lineNumber":99,"sourceCode":"    return f'Zero({self.aval})'\n  def instantiate(self):\n    return zeros_like_aval(self.aval)\n\nregister_pytree_node(Zero, lambda z: ((), z.aval), lambda aval, _: Zero(aval))\n\ndef p2tz(primal_value):\n  return Zero(typeof(primal_value).to_tangent_aval())\n\ndef p2cz(primal_value):\n  return Zero(typeof(primal_value).to_ct_aval())\n\ndef a2tz(primal_aval):\n  return Zero(primal_aval.to_tangent_aval())\n\n\ndef _stop_gradient_impl[T](x: T) -> T:\n  if not core.valid_jaxtype(x):\n    raise TypeError(\"stop_gradient only works on valid JAX arrays, but \"\n                    f\"input argument is: {x}\")\n  return x\n\nstop_gradient_p : Primitive = Primitive('stop_gradient')\nstop_gradient_p.def_impl(_stop_gradient_impl)\nstop_gradient_p.def_abstract_eval(lambda x: x)\n\n\n# User-facing version of `Zero`\nclass SymbolicZero:\n  def __init__(self, aval: core.AbstractValue) -> None:\n    self.aval = aval\n\n  def __repr__(self) -> str:\n    return self.__class__.__name__\n\n  # TODO(mattjj,frostig): this forwards attr lookup to self.aval delegate;\n  # should dedup with core.Tracer.__getattr__ which does the same thing","sourceCodeStart":81,"sourceCodeEnd":117,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ad_util.py#L81-L117","documentation":"FromPickle found a node tuple whose size is not the expected 6 elements. A pickled PyTreeDef stream is corrupt or was produced by an incompatible serialization format/version.","triggerScenarios":"Unpickling a PyTreeDef whose __getstate__ payload contains tuples of length != 6, e.g. hand-crafted pickles or pickles from a different jaxlib version with a different node format.","commonSituations":"Loading checkpoints/saved states across jaxlib versions; truncated or manually edited pickle data; pickle protocol incompatibilities.","solutions":["Re-save the pickle with the currently running jaxlib version","Verify the pickle payload integrity (checksum/size) before loading","Avoid hand-editing pickled treedefs; recompute structures from data"],"exampleFix":"# before\ntreedef = pickle.loads(old_checkpoint['treedef'])  # from older jaxlib\n# after\n# regenerate: _, treedef = jax.tree_util.tree_flatten(fresh_params)","handlingStrategy":"fallback","validationCode":"import jax\n# check pickle payload shape before unpickling where possible\n# simplest: wrap unpickle and fall back to regeneration","typeGuard":null,"tryCatchPattern":"try:\n    treedef = pickle.loads(data)\nexcept Exception as e:\n    if 'Malformed pickled PyTreeDef' in str(e):\n        _, treedef = jax.tree_util.tree_flatten(fresh_reference_obj)\n    else:\n        raise","preventionTips":["Save jaxlib version alongside pickled treedefs","Regenerate structures from data rather than unpickling when versions differ","Checksum checkpoint files"],"tags":["pytree","jax","pickle","version-mismatch"],"backgroundTag":"pytree-pickle-corruption","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}