{"record":{"id":"fc9f9cc5301e433e","repo":"jax-ml/jax","slug":"custom-jvp-rule-jvp-name-for-function-self-trac-fc9f9c","errorCode":null,"errorMessage":"Custom JVP rule {jvp_name} for function {self.traced.fun_name} must produce primal and tangent outputs with equal container (pytree) structures, but got {self.out_tree} and {tree} respectively.","messagePattern":"Custom JVP rule (.+?) for function (.+?) must produce primal and tangent outputs with equal container \\(pytree\\) structures, but got (.+?) and (.+?) respectively\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1054,"sourceCode":"    if self.symbolic_zeros:\n      tangents_ = tree_map(ad_util.replace_internal_symbolic_zeros, tangents_,\n                           is_leaf=zero)\n    else:\n      tangents_ = tree_map(ad_util.instantiate, tangents_, is_leaf=zero)\n    pair_out = self.jvp_fun(*static_args, primals_, tangents_)\n    jvp_name = getattr(self.jvp_fun, '__name__', str(self.jvp_fun))\n    if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:\n      raise TypeError(\n          f\"Custom JVP rule {jvp_name} for function {self.traced.fun_name} \"\n          \"must produce a pair (list or tuple of length two) representing \"\n          f\"primal and tangent outputs, but got {pair_out}.\")\n    out, out_tangent = pair_out\n    if (tree := tracing_registry.flatten(out)[1]) != self.out_tree:\n      raise TypeError(_jvp_primal_tree_mismatch_err(self, jvp_name, out))\n    _jvp_check_primal_avals(self, jvp_name, out)\n    zero_ = lambda x: isinstance(x, (ad_util.Zero, ad_util.SymbolicZero))\n    if (tree := tracing_registry.flatten(out_tangent, zero_)[1]) != self.out_tree:\n      raise TypeError(\n          f\"Custom JVP rule {jvp_name} for function {self.traced.fun_name} \"\n          \"must produce primal and tangent outputs with equal container \"\n          f\"(pytree) structures, but got {self.out_tree} and {tree} \"\n          \"respectively.\")\n    _jvp_check_tangent_avals(self, out, out_tangent)\n    out_tangent = tree_map(ad_util.replace_rule_output_symbolic_zeros,\n                           out_tangent, is_leaf=zero_)\n    return out, out_tangent\n\n  lin, linearized = linearize_from_jvp\n  vjp_fwd, vjp_bwd_retval = vjp_from_jvp\n\n  def transpose(self, out_ct, *args):\n    # The application must be linear in the accumulated args\n    args_flat = tree_leaves_checked(self.in_tree, args)\n    is_lin = [isinstance(x, ad.GradAccum) for x in args_flat]\n    vals = [x for x, l in zip(args_flat, is_lin) if not l]\n","sourceCodeStart":1036,"sourceCodeEnd":1072,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1036-L1072","documentation":"The primal and tangent outputs of a custom JVP rule must have identical pytree structure. JAX flattens the tangent output (treating Zero/SymbolicZero as leaves) and compares its tree to the primal's; any structural difference raises this error showing both trees.","triggerScenarios":"A defjvp whose tangent output has different nesting than the primal output — e.g. primal returns (a, b) but tangent is a single array, or tangent omits a None entry present in the primal.","commonSituations":"Rules that compute tangents only for some outputs and return a partial structure; using jax.tree_util.tree_map on one side but not the other; asymmetry introduced when adding a new output to the function.","solutions":["Return tangent structure mirroring the primal exactly, using ad.InstantiatedZero / jnp.zeros-like values (or ad_util zeros) for non-differentiated outputs","Construct both via tree_map over the same template: primal_out, tangent_out = tree_map(...), tree_map(...) on the same structure","Check with jax.tree_util.tree_structure(out) == tree_structure(out_tangent) in tests"],"exampleFix":"// before\n@f.defjvp\ndef f_jvp(p, t):\n  (x,), (xd,) = p, t\n  return (x * 2, x + 1), 2 * xd   # tangent not a pair\n// after\n@f.defjvp\ndef f_jvp(p, t):\n  (x,), (xd,) = p, t\n  return (x * 2, x + 1), (2 * xd, xd)","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\nout, out_t = f_jvp(primals, tangents)\nassert jtu.tree_structure(out) == jtu.tree_structure(out_t)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use zeros/None consistently per API for non-differentiated outputs","Construct both outputs with the same tree_map template"],"tags":["jax","custom-jvp","pytree-mismatch"],"backgroundTag":"jax-custom-jvp-tangent-pytree-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}