{"record":{"id":"c99bab940672ea6b","repo":"jax-ml/jax","slug":"tree-mismatch-during-linearization-of-prim-exp","errorCode":null,"errorMessage":"tree mismatch during linearization of {prim=}. Expected: {prim.out_tree} got: {treedef_actual}","messagePattern":"tree mismatch during linearization of (.+?)\\. Expected: (.+?) got: (.+?)","errorType":"error_code","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":456,"sourceCode":"  nz_tangents = tree_leaves(tangents)\n  out_nz = call_hi_primitive_linearized_p.bind(\n      *residuals_flat, *nz_tangents, residuals_tree=residuals_tree, _prim=prim,\n      nz_in_flat=tuple(nz_in_flat), nz_out_flat=tuple(nz_out_flat),\n      has_sres=sres is not None)\n  out_nz_iter = iter(out_nz)\n  out = [next(out_nz_iter) if nz else ad_util.Zero(a.to_tangent_aval())\n         for a, nz in zip(prim.out_avals_flat, nz_out_flat)]\n  assert next(out_nz_iter, sentinel := object()) is sentinel\n  return out\n\ndef flatten_user_linearized(prim, residuals, sres, *tangents_flat):\n  tangents = tree_unflatten(prim.in_tree, tangents_flat)\n  tangents_out = (prim.linearized(residuals, *tangents) if sres is None else\n                  prim.linearized(residuals, sres, *tangents))\n  flat_vals, treedef_actual = tracing_registry.flatten(\n      tangents_out, lambda x: isinstance(x, ad_util.Zero))\n  if treedef_actual != prim.out_tree:\n    raise RuntimeError(\n        f\"tree mismatch during linearization of {prim=}.\"\n        f\" Expected: {prim.out_tree} got: {treedef_actual}\"\n    )\n  return flat_vals\n\ncall_hi_primitive_linearized_p = core.Primitive(\"call_hi_primitive_linearized\")\ncall_hi_primitive_linearized_p.multiple_results = True\ncall_hi_primitive_linearized_p.is_high = lambda *args, _prim, **_: True\n@call_hi_primitive_linearized_p.def_abstract_eval\ndef _call_hi_primitive_linearized_abstract_eval(\n    *_args, _prim, residuals_tree, nz_in_flat, nz_out_flat, has_sres):\n  return [t.to_tangent_aval() for t, nz in zip(_prim.out_avals_flat, nz_out_flat) if nz]\n\ndef _call_hi_primitive_linearized_transpose(\n    cts_flat_, *args, _prim, residuals_tree, nz_in_flat, nz_out_flat, has_sres):\n  residuals_flat, accums_flat = split_list(args, [residuals_tree.num_leaves])\n  residuals = tree_unflatten(residuals_tree, residuals_flat)\n  accums_flat_ = iter(accums_flat)","sourceCodeStart":438,"sourceCodeEnd":474,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L438-L474","documentation":"The user-supplied `linearized` rule of a HiPrim returned tangents whose pytree structure doesn't match the primitive's declared out_tree. flatten_user_linearized flattens the rule output and compares treedefs, raising RuntimeError on mismatch.","triggerScenarios":"A custom linearized(residuals, *tangents) implementation returning e.g. a tuple instead of the dict, wrong arity, or differently-nested containers than the primal output tree.","commonSituations":"Hand-written linearize rules that forget an output component or wrap results in an extra list/tuple.","solutions":["Make linearized return exactly the same pytree structure as the primal output (mirror out_tree)","Use jax.tree_util.tree_unflatten(prim.out_tree, ...) or build the output the same way the primal does","Add a unit test comparing tree_structure(prim.linearized(...)) with tree_structure(primal output)"],"exampleFix":"# before\ndef linearized(self, residuals, *ts):\n    return [ts[0]]  # wrong structure\n# after\ndef linearized(self, residuals, *ts):\n    return {'y': ts[0]}  # matches out_tree dict","handlingStrategy":"type-guard","validationCode":"from jax.tree_util import tree_structure\nassert tree_structure(dummy_out) == tree_structure(prim.out_tree.unflatten([0]*n))","typeGuard":"def linearized_tree_ok(prim, residuals, tangents) -> bool:\n    from jax.tree_util import tree_structure\n    return tree_structure(prim.linearized(residuals, *tangents)) == prim.out_tree","tryCatchPattern":"try:\n    f_lin(t)\nexcept RuntimeError as e:\n    if 'tree mismatch during linearization' in str(e):\n        raise RuntimeError('fix linearized return structure to match primal out_tree') from e\n    raise","preventionTips":["Construct rule outputs with the same container shapes as the primal","Unit-test tree_structure equality between primal and linearized outputs"],"tags":["jax","pytree","linearize","autodiff"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}