{"record":{"id":"7a70ddde6397e6d7","repo":"jax-ml/jax","slug":"primal-and-tangent-arguments-to-jax-jvp-do-not-mat","errorCode":null,"errorMessage":"primal and tangent arguments to jax.jvp do not match; dtypes must be equal, or in case of int/bool primal dtype the tangent dtype must be float0.Got primal dtype {_dtype(p)} and so expected tangent dtype {core.primal_dtype_to_tangent_dtype(_dtype(p))}, but got tangent dtype {_dtype(t)} instead.","messagePattern":"primal and tangent arguments to jax\\.jvp do not match; dtypes must be equal, or in case of int/bool primal dtype the tangent dtype must be float0\\.Got primal dtype (.+?) and so expected tangent dtype (.+?), but got tangent dtype (.+?) instead\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1462,"sourceCode":"  \"\"\"\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\n@overload\ndef linearize(fun: Callable, *primals, has_aux: Literal[False] = False,\n              in_nzs: Any = None) -> tuple[Any, Callable]:\n  ...\n\n@overload","sourceCodeStart":1444,"sourceCodeEnd":1480,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1444-L1480","documentation":"Each tangent leaf of jax.jvp must have the dtype matching the primal's tangent dtype: identical for float primals, and float0 for int/bool primals. This error fires when, e.g., an int primal gets a float32 tangent or a float primal gets a mismatched float dtype.","triggerScenarios":"jax.jvp(f, (n,), (jnp.ones(3),)) where n is an int array (tangent must be float0); a f32 primal paired with an f64 tangent; tangents created with jnp.zeros instead of jnp.zeros_like.","commonSituations":"Index arguments treated as differentiable; mixed-precision (f32/f64) code with enable_x64 inconsistencies; ones_like applied to a Python int instead of the array.","solutions":["Create tangents with jax.tree.map(jnp.zeros_like, primals) or jax.ad.instantiate_zeros to get correct dtypes including float0","Cast tangents to match primal dtype (or float0 for int/bool primals)","Ensure enable_x64 setting matches between primal creation and tangent creation"],"exampleFix":"# before\njax.jvp(f, (idx,), (jnp.ones_like(idx),))  # idx is int\n# after\nfrom jax.ad import instantiate_zeros\njax.jvp(f, (idx,), instantiate_zeros((idx,)))","handlingStrategy":"validation","validationCode":"for p, t in zip(tree_leaves(primals), tree_leaves(tangents)):\n    if hasattr(p, 'dtype'):\n        assert core.primal_dtype_to_tangent_dtype(p.dtype) == t.dtype, f'tangent dtype mismatch for {p.dtype} -> {t.dtype}'","typeGuard":"def tangent_dtype_ok(p, t):\n    return not hasattr(p, 'dtype') or jax.core.primal_dtype_to_tangent_dtype(p.dtype) == t.dtype","tryCatchPattern":null,"preventionTips":["Use jax.ad.instantiate_zeros or tree.map(zeros_like) for tangents","Remember int/bool primals need float0 tangents","Keep enable_x64 consistent across the program"],"tags":["jax","jvp","dtype","autodiff","float0"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}