{"record":{"id":"bbe452cbe5952ce9","repo":"jax-ml/jax","slug":"custom-jvp-rule-must-produce-primal-and-tangent-ou-bbe452","errorCode":null,"errorMessage":"Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes. Expected {} (tangent type of {}) but got {}.","messagePattern":"Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes\\. Expected (.+?) \\(tangent type of (.+?)\\) but got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1152,"sourceCode":"\ndef _jvp_check_tangent_avals(self, out, out_tangent):\n  strip = lambda a: a.strip_weak_type() if hasattr(a, 'strip_weak_type') else a\n  out_flat = tree_leaves_checked(self.out_tree, out)\n  tangents_flat = self.out_tree.flatten_up_to(out_tangent)\n  primal_avals_out = [strip(typeof(x)) for x in out_flat]\n  expected_tangent_avals_out = [a.to_tangent_aval() for a in primal_avals_out]\n  tangent_avals_out = [\n      strip(t.aval) if isinstance(t, (ad_util.Zero, ad_util.SymbolicZero))\n      else strip(typeof(t)) for t in tangents_flat]\n  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","sourceCodeStart":1134,"sourceCodeEnd":1170,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1134-L1170","documentation":"For a custom_jvp function with a single output, the tangent output of the JVP rule must be the tangent-type of the primal output (same shape; float dtype if the primal is float; float0 if primal is non-differentiable). This TypeError reports the expected vs actual tangent aval.","triggerScenarios":"Returning a tangent with a different shape than the primal (e.g. forgetting to broadcast), an integer tangent for a float primal, or a scalar tangent for an array primal.","commonSituations":"Rules that return jnp.ones_like(primal) with a cast, integer-valued derivative expressions (e.g. from integer constants), tangent computed at wrong shapes when primal gets broadcast.","solutions":["Match tangent shape/dtype to the primal: use jnp.ones_like / result * ones_like(primal) rather than new arrays with other dtypes","Wrap derivative constants as primal-dtype: 2.0 * xd or jnp.array(2.0, xd.dtype)","Test with jax.jvp(f, (x,), (t,)) for representative inputs"],"exampleFix":"// before\n@f.defjvp\ndef f_jvp(p, t):\n  (x,), (xd,) = p, t\n  return x * 2, 2        # scalar int tangent\n// after\n@f.defjvp\ndef f_jvp(p, t):\n  (x,), (xd,) = p, t\n  return x * 2, 2 * xd  # array tangent matching primal","handlingStrategy":"validation","validationCode":"av_p, av_t = jax.eval_shape(lambda p, t: (f_jvp(p, t)[0], f_jvp(p, t)[1]), (x,), (t,))\nassert av_p.shape == av_t.shape and av_t.dtype == jnp.result_type(av_p.dtype, jnp.float32) or av_p.dtype == av_t.dtype","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Return 2.0 * tangent-style arrays rather than bare scalars","Test with jax.jvp on representative inputs"],"tags":["jax","custom-jvp","tangent-mismatch"],"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"}