{"record":{"id":"5c31dc1c78c8692e","repo":"jax-ml/jax","slug":"custom-jvp-rule-jvp-name-for-function-primal-na-5c31dc","errorCode":null,"errorMessage":"Custom JVP rule {jvp_name} for function {primal_name} must produce primal and tangent outputs with equal container (pytree) structures, but got {out_tree} and {out_tree2} respectively.","messagePattern":"Custom JVP rule (.+?) for function (.+?) must produce primal and tangent outputs with equal container \\(pytree\\) structures, but got (.+?) and (.+?) respectively\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":323,"sourceCode":"def _flatten_jvp(f, store, primal_name, jvp_name, in_tree, maybe_out_type, *args):\n  primals_in, tangents_in = split_list(args, [len(args) // 2])\n  py_primals = tree_unflatten(in_tree, primals_in)\n  py_tangents = tree_unflatten(in_tree, tangents_in)\n  pair_out = f(py_primals, py_tangents)\n  if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:\n    msg = (f\"Custom JVP rule {jvp_name} for function {primal_name} \"\n           \"must produce a pair (list or tuple of length two) representing \"\n           f\"primal and tangent outputs, but got {pair_out}.\")\n    raise TypeError(msg)\n  py_primals_out, py_tangents_out = pair_out\n  primals_out, out_tree = tree_flatten(py_primals_out)\n  tangents_out, out_tree2 = tree_flatten(py_tangents_out)\n  primal_avals = [core.typeof(x) for x in primals_out]\n  if out_tree != out_tree2:\n    msg = (f\"Custom JVP rule {jvp_name} for function {primal_name} must \"\n           \"produce primal and tangent outputs with equal container (pytree) \"\n           f\"structures, but got {out_tree} and {out_tree2} respectively.\")\n    raise TypeError(msg)\n  # If the primal function already ran, check out_tree agreement.\n  try: out_type_ = maybe_out_type()\n  except lu.StoreException: out_type_ = None\n  if out_type_ is not None:\n    out_tree_, primal_avals_, () = out_type_\n    ty_tree  = tree_unflatten(out_tree , [a.str_short() for a in primal_avals])\n    ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_])\n    if out_tree_ != out_tree:\n      m = (f\"Custom JVP rule {jvp_name} for function {primal_name} must \"\n           \"produce a pair (list or tuple of length two) \"\n           \"where the first element represents the primal output \"\n           \"(equal in value to the output of the custom_jvp-decorated function \"\n           f\"{primal_name}, \"\n           \"and in particular of the same container/pytree structure), but \"\n           \"instead the JVP rule output's first element had container/pytree \"\n           \"structure:\\n\"\n           f\"\"\"    {str(ty_tree ).replace(\"'\", \"\")}\\n\"\"\"\n           f\"while the custom_jvp-decorated function {primal_name} had output \"","sourceCodeStart":305,"sourceCodeEnd":341,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L305-L341","documentation":"The primal and tangent outputs returned by a custom_jvp rule must have identical pytree container structure. This fires when, e.g., the primal is a tuple of two arrays but the tangent is a single array or a differently nested tuple.","triggerScenarios":"A defjvp rule whose primal output is (a, b) but whose tangent is a scalar zeros; mixing list/tuple nesting between the two halves of the returned pair.","commonSituations":"Multi-output functions where only one output's tangent was implemented; constructing tangents with jnp.zeros_like applied to the wrong container.","solutions":["Build the tangent as the same pytree as the primal, e.g. tree_map(jnp.zeros_like, primal_out) then fill in real tangents","Mirror tuple nesting exactly between the two returned values"],"exampleFix":"# before\nreturn (a, b), dz\n# after\nreturn (a, b), (dz, jnp.zeros_like(b))","handlingStrategy":"validation","validationCode":"import jax\nassert jax.tree_util.tree_structure(primal_out) == jax.tree_util.tree_structure(tangent_out)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Construct tangents via tree_map(jnp.zeros_like, primal_out)","Keep primal/tangent packing symmetric in rule code"],"tags":["jax","custom-jvp","pytree","autodiff"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}