{"record":{"id":"fec1bddeec04365b","repo":"jax-ml/jax","slug":"custom-jvp-rule-jvp-name-for-function-self-trac-fec1bd","errorCode":null,"errorMessage":"Custom JVP rule {jvp_name} for function {self.traced.fun_name} must produce a pair (list or tuple of length two) where the first element represents the primal output (equal in value to the output of the custom_jvp-decorated function {self.traced.fun_name}, and in particular with leaves of the same shape/dtype), but instead the JVP rule output's first element had shapes/dtypes of:\n    {str(ty_tree ).replace(\"'\", \"\")}\nwhile the custom_jvp-decorated function {self.traced.fun_name} had output shapes/dtypes of:\n    {str(ty_tree_).replace(\"'\", \"\")}","messagePattern":"Custom JVP rule (.+?) for function (.+?) must produce a pair \\(list or tuple of length two\\) where the first element represents the primal output \\(equal in value to the output of the custom_jvp-decorated function (.+?), and in particular with leaves of the same shape/dtype\\), but instead the JVP rule output's first element had shapes/dtypes of:\n    (.+?)\nwhile the custom_jvp-decorated function (.+?) had output shapes/dtypes of:\n    (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1050,"sourceCode":"    primals_ = tuple(x for x in primals if not isinstance(x, Static))\n    tangents_ = tuple(t for x, t in zip(primals, tangents)\n                      if not isinstance(x, Static))\n    zero = lambda x: isinstance(x, ad_util.Zero)\n    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","sourceCodeStart":1032,"sourceCodeEnd":1068,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1032-L1068","documentation":"The first element (primal output) of the custom JVP rule's return pair has a different pytree structure than the output of the @jax.custom_jvp-decorated function. JAX flattens both and compares trees; a mismatch (extra nesting, dict vs tuple, None placement) triggers this error with a detailed diff of shapes/dtypes.","triggerScenarios":"A defjvp returning e.g. a flat array where the decorated function returns (array,), or returning a dict where the function returns a tuple, or omitting a None output.","commonSituations":"Decorated function returns multiple outputs but the JVP rule returns one packed object; refactoring output structure without updating the rule; mixing tuple/list conventions inconsistently across the pair.","solutions":["Make the JVP rule's first element structurally identical to the decorated function's return (same tuple nesting, keys, Nones)","Wrap/unpack outputs symmetrically: if f returns (a, b), the rule must return (a, b), tan_a, tan_b as (primal_pair, tangent_pair)","Add a quick unit test comparing jax.tree.structure(f(x)) with the rule's first output"],"exampleFix":"// before\n@jax.custom_jvp\ndef f(x):\n  return x * 2, x + 1\n@f.defjvp\ndef f_jvp(p, t):\n  (x,), (xd,) = p, t\n  return x * 2, 2 * xd        # primal is single array, should be 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\nstruct_f = jtu.tree_structure(jax.eval_shape(f, x))\nstruct_rule = jtu.tree_structure(jax.eval_shape(lambda p, t: f_jvp(p, t)[0], (x,), (t,)))\nassert struct_f == struct_rule","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build primal and tangent halves of the rule's return symmetrically","Keep a structural unit test comparing trees"],"tags":["jax","custom-jvp","pytree-mismatch"],"backgroundTag":"jax-custom-jvp-primal-pytree-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}