{"record":{"id":"399ab5dc08f88cb8","repo":"jax-ml/jax","slug":"custom-jvp-rule-jvp-name-for-function-primal-na-399ab5","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 of the same container/pytree structure), but instead the JVP rule output's first element had container/pytree structure:\\n    {str(ty_tree ).replace(\"'\", \"\")}\\nwhile the custom_jvp-decorated function {primal_name} had output container/pytree structure:\\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 of the same container/pytree structure\\), but instead the JVP rule output's first element had container/pytree structure:\\\\n    (.+?)\\\\nwhile the custom_jvp-decorated function (.+?) had output container/pytree structure:\\\\n    (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":344,"sourceCode":"  except lu.StoreException: out_type_ = None\n  if out_type_ is not None:\n    out_tree_, primal_avals_, () = out_type_\n    ty_tree  = tree_unflatten(out_tree , [a.str_short() for a in primal_avals])\n    ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_])\n    if out_tree_ != out_tree:\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 of the same container/pytree structure), but \"\n           \"instead the JVP rule output's first element had container/pytree \"\n           \"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()","sourceCodeStart":326,"sourceCodeEnd":362,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L326-L362","documentation":"The first element of a custom_jvp rule's returned pair must match the container/pytree structure of the original function's output (which JAX knows from a prior primal trace). This error reports the two differing structures side by side.","triggerScenarios":"The rule returns the primal output wrapped differently than the original function does — e.g. original returns a scalar, rule returns a 1-tuple, or nesting levels differ.","commonSituations":"Rules written before the primal function was refactored to return a different container; generic rule code that always wraps outputs in lists.","solutions":["Make the rule's first returned element structurally identical to what the undecorated function returns","Run the function once without grad to inspect its output structure and replicate it"],"exampleFix":"# before\nreturn [y], dy\n# after\nreturn y, dy","handlingStrategy":"validation","validationCode":"import jax\norig_tree = jax.tree_util.tree_structure(f(*sample_args))\n# assert rule's first return matches orig_tree before returning","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Snapshot the function's output structure in tests","Avoid defensive wrapping/unwrapping in rules"],"tags":["jax","custom-jvp","pytree","autodiff"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}