{"record":{"id":"960885ae5b26195c","repo":"jax-ml/jax","slug":"custom-jvp-rule-jvp-name-for-function-primal-na-960885","errorCode":null,"errorMessage":"Custom JVP rule {jvp_name} for function {primal_name} must produce a pair (list or tuple of length two) where the first element represents the primal output (equal in value to the output of the custom_jvp-decorated function {primal_name}, and in particular with leaves of the same shape/dtype), but instead the JVP rule output's first element had shapes/dtypes of:\\n    {str(ty_tree ).replace(\"'\", \"\")}\\nwhile the custom_jvp-decorated function {primal_name} had output shapes/dtypes of:\\n    {str(ty_tree_).replace(\"'\", \"\")}","messagePattern":"Custom JVP rule (.+?) for function (.+?) must produce a pair \\(list or tuple of length two\\) where the first element represents the primal output \\(equal in value to the output of the custom_jvp-decorated function (.+?), and in particular with leaves of the same shape/dtype\\), but instead the JVP rule output's first element had shapes/dtypes of:\\\\n    (.+?)\\\\nwhile the custom_jvp-decorated function (.+?) had output shapes/dtypes of:\\\\n    (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":357,"sourceCode":"           \"structure:\\n\"\n           f\"\"\"    {str(ty_tree ).replace(\"'\", \"\")}\\n\"\"\"\n           f\"while the custom_jvp-decorated function {primal_name} had output \"\n           \"container/pytree structure:\\n\"\n           f\"\"\"    {str(ty_tree_).replace(\"'\", \"\")}.\"\"\")\n      raise TypeError(m)\n    if not all(map(core.typematch, primal_avals, primal_avals_)):\n      m = (f\"Custom JVP rule {jvp_name} for function {primal_name} must \"\n           \"produce a pair (list or tuple of length two) \"\n           \"where the first element represents the primal output \"\n           \"(equal in value to the output of the custom_jvp-decorated function \"\n           f\"{primal_name}, \"\n           \"and in particular with leaves of the same shape/dtype), but \"\n           \"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}\"","sourceCodeStart":339,"sourceCodeEnd":375,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L339-L375","documentation":"Stronger than the structure check: the leaves of a custom_jvp rule's primal output must have the same shape and dtype as the original function's outputs. This fires when structures match but an array's shape or dtype differs (e.g. broadcasting added a dim, or the rule computed in float64).","triggerScenarios":"A rule that returns y[None] instead of y, squeezes/expands dims, or promotes dtypes (e.g. multiplying by a Python float under x64 mode) relative to the primal output.","commonSituations":"x64 enabled causing silent dtype promotion; reshape/broadcast differences between the rule and the original implementation; stale rule after the primal's dtype contract changed.","solutions":["Return arrays with exactly the primal's shape and dtype; cast with .astype(primal.dtype) if needed","Disable jax_enable_x64 or ensure constants are typed (jnp.float32(...)) to avoid promotion","Compare jax.eval_shape of the function and the rule's primal half"],"exampleFix":"# before\nreturn y * 1.0, dy\n# after\nreturn y.astype(primal_dtype), dy","handlingStrategy":"validation","validationCode":"import jax\npspec = jax.eval_shape(f, *sample_args)\n# in the rule: return jax.tree_util.tree_map(lambda a, b: b.astype(a.dtype), pspec, primal_out), tangent","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Type literals explicitly (jnp.float32(1.0)) inside rules","Watch for x64 promotion in FP8 pipelines"],"tags":["jax","custom-jvp","shape","dtype","autodiff"],"backgroundTag":"shape-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}