{"record":{"id":"3c1c330fed755f7b","repo":"jax-ml/jax","slug":"primal-and-tangent-arguments-to-jax-jvp-must-have","errorCode":null,"errorMessage":"primal and tangent arguments to jax.jvp must have the same tree structure; primals have tree structure {ps_ft.tree} whereas tangents have tree structure {ts_ft.tree}.","messagePattern":"primal and tangent arguments to jax\\.jvp must have the same tree structure; primals have tree structure (.+?) whereas tangents have tree structure (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1456,"sourceCode":"  >>>\n  >>> primals, tangents = jax.jvp(jax.numpy.sin, (0.1,), (0.2,))\n  >>> print(primals)\n  0.09983342\n  >>> print(tangents)\n  0.19900084\n  \"\"\"\n  check_callable(fun)\n  if (not isinstance(primals, (tuple, list)) or\n      not isinstance(tangents, (tuple, list))):\n    raise TypeError(\"primal and tangent arguments to jax.jvp must be tuples or lists; \"\n                    f\"found {type(primals).__name__} and {type(tangents).__name__}.\")\n  return _jvp(fun, primals, tangents, has_aux=has_aux)\n\ndef _jvp(fun: Callable, primals, tangents, has_aux=False):\n  ps_ft = ft.flatten(primals)\n  ts_ft = ft.flatten(tangents)\n  if ps_ft.tree != ts_ft.tree:\n    raise TypeError(\"primal and tangent arguments to jax.jvp must have the same tree \"\n                    f\"structure; primals have tree structure {ps_ft.tree} whereas tangents have \"\n                    f\"tree structure {ts_ft.tree}.\")\n  for p, t in zip(ps_ft, ts_ft):\n    if not isinstance(core.typeof(p), ShapedArray): continue\n    if core.primal_dtype_to_tangent_dtype(_dtype(p)) != _dtype(t):\n      raise TypeError(\"primal and tangent arguments to jax.jvp do not match; \"\n                      \"dtypes must be equal, or in case of int/bool primal dtype \"\n                      \"the tangent dtype must be float0.\"\n                      f\"Got primal dtype {_dtype(p)} and so expected tangent dtype \"\n                      f\"{core.primal_dtype_to_tangent_dtype(_dtype(p))}, but got \"\n                      f\"tangent dtype {_dtype(t)} instead.\")\n    if np.shape(p) != np.shape(t):\n      raise ValueError(\"jvp called with different primal and tangent shapes;\"\n                       f\"Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}\")\n\n  out_primals, out_tangents, *aux = ad.jvp(fun, ps_ft, ts_ft, has_aux=has_aux)\n  return out_primals.unflatten(), out_tangents.unflatten(), *aux\n","sourceCodeStart":1438,"sourceCodeEnd":1474,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1438-L1474","documentation":"jax.jvp flattens primals and tangents as pytrees and requires identical tree structures (same containers, keys, arities). This error reports the two pytree definitions when they differ.","triggerScenarios":"jax.jvp(f, {'a': x}, (x,)); primals as a dict but tangents as a tuple; primals with two leaves and tangents with one; differing dict keys.","commonSituations":"Building tangents with a different helper than primals (e.g. jax.tree.map with a different structure, or tree_map with is_leaf inconsistency); adding/removing arguments in one place only.","solutions":["Construct tangents from primals with jax.tree.map(jnp.ones_like, primals) or jax.lax.zeros_like_pytree so structures always match","Fix the container types/keys so both sides match exactly","Add asserts comparing tree structures before calling jvp"],"exampleFix":"# before\njax.jvp(f, (x, y), (t_x,))\n# after\ntangents = jax.tree.map(jnp.zeros_like, (x, y))\njax.jvp(f, (x, y), tangents)","handlingStrategy":"validation","validationCode":"assert jax.tree.structure(primals) == jax.tree.structure(tangents), 'primal/tangent tree mismatch'","typeGuard":"def trees_match(p, t): return jax.tree.structure(p) == jax.tree.structure(t)","tryCatchPattern":null,"preventionTips":["Derive tangents from primals via jax.tree.map(jnp.zeros_like, primals)","Never hand-build tangent containers","Compare tree structures in debug wrappers"],"tags":["jax","jvp","pytree","autodiff"],"backgroundTag":"structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}