{"record":{"id":"9f42a19dd31021fb","repo":"jax-ml/jax","slug":"custom-jvp-rule-must-produce-primal-and-tangent-ou-9f42a1","errorCode":null,"errorMessage":"Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:\\n{}","messagePattern":"Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:\\\\n(.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":378,"sourceCode":"    core.typeof(x).strip_weak_type().to_tangent_aval()\n    for x in primals_out]\n  tangent_avals_out = [core.typeof(t).strip_weak_type()\n                       if type(t) is not SymbolicZero else t.aval.strip_weak_type()\n                       for t in tangents_out]\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, expected_tangent_avals_out, tangent_avals_out\n      msg = (\"Custom JVP rule must produce primal and tangent outputs with \"\n             \"corresponding shapes and dtypes. Expected {} (tangent type of {}) but got {}.\")\n      raise TypeError(msg.format(av_et.str_short(), av_p.str_short(), av_t.str_short()))\n    else:\n      msg = (\"Custom JVP rule must produce primal and tangent outputs with \"\n             \"corresponding shapes and dtypes, but got:\\n{}\")\n      disagreements = (\n          f\"  primal {av_p.str_short()} with tangent {av_t.str_short()}, expecting tangent {av_et}\"\n          for av_p, av_et, av_t in zip(primal_avals_out, expected_tangent_avals_out, tangent_avals_out)\n          if av_et != av_t)\n      raise TypeError(msg.format('\\n'.join(disagreements)))\n  store.store((out_tree, primal_avals, ()))\n  return primals_out + tangents_out\n\nclass CustomJVPCallPrimitive(core.Primitive):\n  multiple_results = True\n  skip_canonicalization = True\n\n  def bind_with_trace(self, trace, args, avals, params, /):\n    params = dict(params)\n    fun, jvp = params.pop('subfuns')\n    return trace.process_custom_jvp_call(self, fun, jvp, args, **params)\n\n  def impl(self, fun, _, *args):\n    raise NotImplementedError\n\n  def is_high(self, *_, call_jaxpr, **__):\n    return call_jaxpr.is_high\n","sourceCodeStart":360,"sourceCodeEnd":396,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L360-L396","documentation":"Multi-output version of the tangent shape/dtype check: for at least one output leaf, the returned tangent does not match the tangent type of the corresponding primal output; the message lists each disagreement.","triggerScenarios":"A custom_jvp rule over a multi-output function where one or more tangents have wrong shape/dtype — e.g. tangent pytree built per-input rather than per-output.","commonSituations":"Functions returning (value, aux) where aux is integer/non-differentiable and the rule returns a float tangent for it; partially updated rules after adding a new output.","solutions":["Build tangents leaf-by-leaf from the primal outputs (tree_map over the primal pytree)","For non-differentiable outputs, ensure the tangent matches their tangent type (often zeros of the corresponding dtype)","Run jax.eval_shape on the rule to inspect every output leaf"],"exampleFix":"# before\nreturn (y, aux), (dy, d_aux_float)\n# after\nreturn (y, aux), (dy, zeros_tangent_of(aux))","handlingStrategy":"validation","validationCode":"import jax\nok = all(a.shape == t.shape and a.dtype == t.dtype\n         for a, t in zip(jax.tree_util.tree_leaves(primal_out),\n                         jax.tree_util.tree_leaves(tangent_out)))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build multi-output tangents with tree_map over the primal pytree","Handle non-differentiable aux outputs with matching zero tangents"],"tags":["jax","custom-jvp","tangent","pytree","autodiff"],"backgroundTag":"shape-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}