{"record":{"id":"1a23a7df3cdc6423","repo":"jax-ml/jax","slug":"custom-jvp-rule-must-produce-primal-and-tangent-ou-1a23a7","errorCode":null,"errorMessage":"Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:\n{disagreements}","messagePattern":"Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:\n(.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1162,"sourceCode":"  if not all(map(core.typematch, expected_tangent_avals_out, tangent_avals_out)):\n    if len(expected_tangent_avals_out) == 1:\n      (av_p,), (av_et,), (av_t,) = (primal_avals_out,\n                                    expected_tangent_avals_out,\n                                    tangent_avals_out)\n      msg = (\"Custom JVP rule must produce primal and tangent outputs with \"\n             \"corresponding shapes and dtypes. \"\n             \"Expected {} (tangent type of {}) but got {}.\")\n      raise TypeError(msg.format(av_et.str_short(), av_p.str_short(),\n                                 av_t.str_short()))\n    else:\n      disagreements = \"\\n\".join(\n          f\"  primal {av_p.str_short()} with tangent {av_t.str_short()}, \"\n          f\"expecting tangent {av_et}\"\n          for av_p, av_et, av_t in zip(primal_avals_out,\n                                       expected_tangent_avals_out,\n                                       tangent_avals_out)\n          if not core.typematch(av_et, av_t))\n      raise TypeError(\n          \"Custom JVP rule must produce primal and tangent outputs with \"\n          f\"corresponding shapes and dtypes, but got:\\n{disagreements}\")\n\n\nclass custom_jvp3:\n  jvp_fun: Callable | None = None\n  symz: bool = False\n\n  def __init__(self, f, nondiff_argnums=(), nondiff_argnames=()):\n    self.static_argnums = _set_up_nondiff(f, nondiff_argnums, nondiff_argnames)\n    update_wrapper(self, f)\n    self.f = f\n\n  def defjvp(self, jvp, symbolic_zeros=False):\n    self.jvp_fun = jvp\n    self.symz = symbolic_zeros\n    return jvp\n","sourceCodeStart":1144,"sourceCodeEnd":1180,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1144-L1180","documentation":"Multi-output variant of the tangent check: JAX zips primal, expected-tangent, and actual-tangent avals and lists every disagreement where the rule's tangent does not typematch the expected tangent type of the primal; the combined message enumerates each mismatched output.","triggerScenarios":"A custom_jvp function returning several outputs where one or more tangents have wrong shape/dtype — e.g. swapping the order of tangents relative to primals, or returning zeros with wrong dtype for one output.","commonSituations":"Tuple outputs whose tangent tuple elements are permuted or truncated; mixed float/int outputs where int outputs need float0/zero tangents; large refactors of multi-output functions.","solutions":["Fix each listed mismatch; the error names primal, its tangent, and the expected tangent per output","Keep tangent construction symmetric with primal construction (build both in the same tree_map)","For non-differentiable outputs return the appropriate zero (e.g. ad.InstantiatedZero / None per API)"],"exampleFix":"// before\n@f.defjvp\ndef f_jvp(p, t):\n  (x,), (xd,) = p, t\n  return (x*2, x+1), (xd, jnp.ones(3))  # second tangent wrong shape\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":"outs, tans = f_jvp(primals, tangents)\njax.tree_util.tree_map(lambda p, t: (p.shape == t.shape), outs, tans)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep multi-output tangents order-aligned with primals","Build both tuples in one place in the code"],"tags":["jax","custom-jvp","tangent-mismatch","multi-output"],"backgroundTag":"jax-custom-jvp-tangent-aval-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}