{"record":{"id":"e2944b85fc7cea0f","repo":"jax-ml/jax","slug":"custom-jvp-rule-jvp-name-for-function-self-trac","errorCode":null,"errorMessage":"Custom JVP rule {jvp_name} for function {self.traced.fun_name} must produce a pair (list or tuple of length two) representing primal and tangent outputs, but got {pair_out}.","messagePattern":"Custom JVP rule (.+?) for function (.+?) must produce a pair \\(list or tuple of length two\\) representing primal and tangent outputs, but got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1044,"sourceCode":"  def expand(self, *args):\n    args = [x for x in args if not isinstance(x, Static)]\n    return self.traced(*args)\n\n  def jvp(self, primals, tangents):\n    static_args = tuple(x.val for x in primals if isinstance(x, Static))\n    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","sourceCodeStart":1026,"sourceCodeEnd":1062,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1026-L1062","documentation":"A custom JVP rule (the function passed to @jax.custom_jvp's defjvp or the jvpfun) must return exactly a length-2 sequence: (primal_out, tangent_out). If it returns anything else (a single array, a 3-tuple, a dict), this TypeError is raised.","triggerScenarios":"Writing a defjvp that returns only the primal output, returns a tuple of 3, returns None, or forgets to also compute and return tangents (e.g. returns out, out_tan, extra_debug).","commonSituations":"First-time custom_jvp users modeling on custom_vjp examples (which have different conventions); refactoring a rule and dropping the tangent return; returning a generator or other non-sequence.","solutions":["Return exactly `return primal_out, tangent_out` from the JVP rule","Ensure both elements are pytrees matching the decorated function's output structure","If tangents are unknown, return zeros of the right shapes (only if semantically correct)"],"exampleFix":"// before\n@jax.custom_jvp\ndef f(x):\n  return x * 2\n@f.defjvp\ndef f_jvp(primals, tangents):\n  (x,), (x_dot,) = primals, tangents\n  return x * 2          # missing tangent\n// after\n@f.defjvp\ndef f_jvp(primals, tangents):\n  (x,), (x_dot,) = primals, tangents\n  return x * 2, 2. * x_dot","handlingStrategy":"validation","validationCode":"out = f_jvp(*static_args, primals_, tangents_)\nassert isinstance(out, (list, tuple)) and len(out) == 2, 'JVP rule must return (primal, tangent)'","typeGuard":"def is_rule_pair(v) -> bool:\n    return isinstance(v, (list, tuple)) and len(v) == 2","tryCatchPattern":null,"preventionTips":["Always end defjvp with `return primals_out, tangents_out`","Unit-test the rule with jax.jvp before integrating"],"tags":["jax","custom-jvp","autodiff","return-contract"],"backgroundTag":"jax-custom-jvp-rule-return-contract","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}