{"record":{"id":"6a5b3b09afee248d","repo":"jax-ml/jax","slug":"but-the-tree-structures-differ","errorCode":null,"errorMessage":"But the tree structures differ:","messagePattern":"But the tree structures differ:","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1819,"sourceCode":"`with_refs` method.\n\nThe arguments to `with_refs` must match the pytree structure of the primal\narguments of the differentiated function {jaxpr.debug_info.func_src_info},\nwith one entry for each array argument. Each entry must be a `Ref` (to\naccumulate this argument's gradient into the ref in-place), a\n`jax.ad.GradValue()` (to have this argument's gradient returned as a value,\nthe default behavior), or a `jax.ad.DontWant()` (to skip computing this\nargument's gradient). Note that `None` is an empty pytree, so it can't be\nused as a placeholder entry.\n\nBut the tree structures differ:\n\"\"\"\n  msg += '\\n'.join(f\"  * args{keystr(path)} was a {thing1} in the primal \"\n                   f\"arguments, but a {thing2} in the `with_refs` arguments, \"\n                   f\"so {explanation}.\"\n                   for path, thing1, thing2, explanation\n                   in equality_errors_pytreedef(in_tree, refs_tree))\n  raise ValueError(msg)\n\ndef _vjp_not_saveable_error(jaxpr, in_tree, idxs):\n  msg = \"\"\"the VJP function was applied before restoring its not-saveable residuals.\n\nBecause `saveable_args` was passed to `jax.vjp`, some argument values that\nwould have been saved for the backward pass were instead replaced with\n`NotSaveable()` sentinels. Before the VJP function can be applied, these\nvalues must be restored, e.g. by assigning to the VJP function's `args_res`\nattribute. The values not yet restored correspond to:\n\"\"\"\n  msg += '\\n'.join(f\"  * {_vjp_arg_name(jaxpr, in_tree, idx)};\" for idx in idxs)\n  raise ValueError(msg)\n\ndef check_accum(aval, acc):\n  if not core.typecompat(acc.aval, aval):\n    raise ValueError(f\"Accumulator aval mismatch: expected {aval}, got {acc.aval}\")\n  return acc\n","sourceCodeStart":1801,"sourceCodeEnd":1837,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1801-L1837","documentation":"When applying a VJP function bound via with_refs, the pytree structure of the refs must match the structure of the original primal arguments. This error enumerates each path where a node was, e.g., a dict in the primal arguments but a tuple in the with_refs arguments, and raises with the itemized diff.","triggerScenarios":"f_vjp.with_refs(grads_tuple)(ct) when primals were passed as a dict; mismatched leaf counts, container types, or dict keys anywhere in the tree.","commonSituations":"Gradient buffers built by a different utility than the parameters (e.g. optax state tree vs param tree); refactoring parameter dataclass fields; mixing namedtuple vs tuple.","solutions":["Build the refs from the primal tree directly: jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primals)","Fix container types/keys to mirror the primal arguments exactly","Log jax.tree.structure(primals) vs jax.tree.structure(refs) to diff structures"],"exampleFix":"# before\nprimals = {'w': w, 'b': b}\nf_vjp.with_refs((g_w, g_b))(ct)\n# after\nprimals = {'w': w, 'b': b}\nrefs = jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primals)\nf_vjp.with_refs(refs)(ct)","handlingStrategy":"validation","validationCode":"assert jax.tree.structure(refs_tree) == jax.tree.structure(primal_tree), 'with_refs tree must mirror primal argument tree'","typeGuard":"def refs_tree_ok(primals, refs): return jax.tree.structure(primals) == jax.tree.structure(refs)","tryCatchPattern":null,"preventionTips":["Build with_refs arguments by tree.map over the primals","Never hand-assemble gradient buffer trees","Diff tree structures in debug logs on failure"],"tags":["jax","vjp","pytree","refs"],"backgroundTag":"structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}