{"record":{"id":"686b8a5d4cf0f0cc","repo":"jax-ml/jax","slug":"custom-jvp-rule-must-produce-primal-and-tangent-ou","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":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":370,"sourceCode":"           \"instead the JVP rule output's first element had shapes/dtypes of:\\n\"\n           f\"\"\"    {str(ty_tree ).replace(\"'\", \"\")}\\n\"\"\"\n           f\"while the custom_jvp-decorated function {primal_name} had output \"\n           \"shapes/dtypes of:\\n\"\n           f\"\"\"    {str(ty_tree_).replace(\"'\", \"\")}\"\"\")\n      raise TypeError(m)\n  primal_avals_out = [core.typeof(x).strip_weak_type() for x in primals_out]\n  expected_tangent_avals_out = [\n    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')","sourceCodeStart":352,"sourceCodeEnd":388,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L352-L388","documentation":"Single-output case of the tangent-check: the tangent returned by a custom_jvp rule must have the shape/dtype equal to the tangent type of the primal output (e.g. float32 tangent for a float32 primal). The message reports expected vs got.","triggerScenarios":"Returning an integer-dtype tangent, a differently shaped tangent (extra/missing batch dim), or a Python scalar where an array is expected, for a single-output custom_jvp function.","commonSituations":"Rules that compute gradients in a different precision; tangent accidentally shaped like the input instead of the output; zeros created with the wrong dtype.","solutions":["Return a tangent with the same shape and dtype as the primal output (use jnp.zeros_like(primal_out) as the base)","Cast/correct explicitly and re-run jax.grad"],"exampleFix":"# before\nreturn y, 0\n# after\nreturn y, jnp.zeros_like(y)","handlingStrategy":"validation","validationCode":"expected = jax.eval_shape(lambda x: jax.zeros(x.shape, x.dtype), primal_out)\n# ensure tangent.shape == primal_out.shape and tangent.dtype == primal_out.dtype before returning","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never return Python scalars as tangents","Base tangents on zeros_like(primal_out)"],"tags":["jax","custom-jvp","tangent","shape","dtype"],"backgroundTag":"shape-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}