{"record":{"id":"a4682e9a75b7eaf8","repo":"jax-ml/jax","slug":"structure-of-the-differentiated-function-jaxpr-de","errorCode":null,"errorMessage":"structure of the differentiated function {jaxpr.debug_info.func_src_info}.\n\nBut the tree structures differ:","messagePattern":"structure of the differentiated function (.+?)\\.\n\nBut the tree structures differ:","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1966,"sourceCode":"\nIf we instead call `f_vjp(2.0, 2.0)`, with the values 'splatted out' as\narguments rather than in a tuple, this error can arise.\n\"\"\".format\n\n\ndef _vjp_ct_tree_error(jaxpr, out_tree, ct_tree):\n  msg = f\"\"\"unexpected tree structure.\n\nThe argument to a VJP function returned by `jax.vjp` must match the pytree\nstructure of the differentiated function {jaxpr.debug_info.func_src_info}.\n\nBut the tree structures differ:\n\"\"\"\n  msg += '\\n'.join(f\"  * out{keystr(path)} was a {thing1} in the original \"\n                   f\"output, but a {thing2} here, so {explanation}.\"\n                   for path, thing1, thing2, explanation\n                   in equality_errors_pytreedef(out_tree, ct_tree))\n  raise ValueError(msg)\n\n\ndef _vjp_check_ct_avals(cts, primal_avals):\n  # TODO(mattjj): improve this error  by flattening with keys in the first place\n  for ct, aval in zip(cts, primal_avals):\n    if isinstance(ct, ad.Zero): continue\n    ct_aval = typeof(ct)\n    ct_aval_expected = aval.to_ct_aval()\n    if (not core.typecompat(ct_aval, ct_aval_expected) and\n        not _temporary_dtype_exception(ct_aval, ct_aval_expected)):\n      raise ValueError(\n          \"unexpected JAX type (e.g. shape/dtype) for argument to VJP function: \"\n          f\"got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} \"\n          \"because the corresponding output of the differentiated function had JAX type \"\n          f\"{aval.str_short()}\")\n\n\n@register_dataclass","sourceCodeStart":1948,"sourceCodeEnd":1984,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1948-L1984","documentation":"In jax.vjp, the cotangent argument you pass to the returned backward function must have the same pytree structure as the primal function's output. JAX detected a structural mismatch (e.g. output was a dict but cotangent is a tuple) and reports exactly which paths differ.","triggerScenarios":"Calling vjp_fn(cts) where cts is a tuple but the differentiated function returned a dict, or a scalar where a pair was returned; changing how outputs are packaged between the forward call and the backward call.","commonSituations":"Function returns multiple values that get auto-packed into a tuple but user passes a single cotangent; using has_aux or auxiliary outputs so the cotangent tree differs from what the user expects; passing jnp.ones(shape) where the output was a pytree of arrays.","solutions":["Pass cotangents matching tree_structure(out) exactly","Recompute out, _ = jax.vjp(f, *args) and mirror the structure of out when building cts","Use jax.tree_util.tree_map(lambda x: jnp.zeros_like(x), out) to build a correctly-structured cotangent"],"exampleFix":"// before\nout, vjp_fn = jax.vjp(f, x)\nin_cts = vjp_fn(1.0)            # f returns (a, b)\n// after\nout, vjp_fn = jax.vjp(f, x)\nin_cts = vjp_fn((jnp.ones_like(a), jnp.ones_like(b)))","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\nout, vjp_fn = jax.vjp(f, *args)\ncts = jtu.tree_map(jnp.zeros_like, out)   # guaranteed structural match\nin_cts = vjp_fn(cts)","typeGuard":null,"tryCatchPattern":"try:\n    in_cts = vjp_fn(cts)\nexcept ValueError as e:\n    if 'tree structures differ' in str(e):\n        cts = jax.tree_util.tree_map(jnp.zeros_like, out)\n        in_cts = vjp_fn(cts)\n    else:\n        raise","preventionTips":["Always derive cotangents from the actual primal output via tree_map","Keep the forward output structure stable across refactors"],"tags":["jax","vjp","autodiff","pytree"],"backgroundTag":"cotangent-tree-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}