jax-ml/jax · error · TypeError

primal and tangent arguments to jax.jvp do not match; dtypes

Error message

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.

What it means

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.

Source

Thrown at jax/_src/api.py:1462

  """
  check_callable(fun)
  if (not isinstance(primals, (tuple, list)) or
      not isinstance(tangents, (tuple, list))):
    raise TypeError("primal and tangent arguments to jax.jvp must be tuples or lists; "
                    f"found {type(primals).__name__} and {type(tangents).__name__}.")
  return _jvp(fun, primals, tangents, has_aux=has_aux)

def _jvp(fun: Callable, primals, tangents, has_aux=False):
  ps_ft = ft.flatten(primals)
  ts_ft = ft.flatten(tangents)
  if ps_ft.tree != ts_ft.tree:
    raise TypeError("primal and tangent arguments to jax.jvp must have the same tree "
                    f"structure; primals have tree structure {ps_ft.tree} whereas tangents have "
                    f"tree structure {ts_ft.tree}.")
  for p, t in zip(ps_ft, ts_ft):
    if not isinstance(core.typeof(p), ShapedArray): continue
    if core.primal_dtype_to_tangent_dtype(_dtype(p)) != _dtype(t):
      raise TypeError("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."
                      f"Got primal dtype {_dtype(p)} and so expected tangent dtype "
                      f"{core.primal_dtype_to_tangent_dtype(_dtype(p))}, but got "
                      f"tangent dtype {_dtype(t)} instead.")
    if np.shape(p) != np.shape(t):
      raise ValueError("jvp called with different primal and tangent shapes;"
                       f"Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}")

  out_primals, out_tangents, *aux = ad.jvp(fun, ps_ft, ts_ft, has_aux=has_aux)
  return out_primals.unflatten(), out_tangents.unflatten(), *aux

@overload
def linearize(fun: Callable, *primals, has_aux: Literal[False] = False,
              in_nzs: Any = None) -> tuple[Any, Callable]:
  ...

@overload

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create tangents with jax.tree.map(jnp.zeros_like, primals) or jax.ad.instantiate_zeros to get correct dtypes including float0
2. Cast tangents to match primal dtype (or float0 for int/bool primals)
3. Ensure enable_x64 setting matches between primal creation and tangent creation

Example fix

# before
jax.jvp(f, (idx,), (jnp.ones_like(idx),))  # idx is int
# after
from jax.ad import instantiate_zeros
jax.jvp(f, (idx,), instantiate_zeros((idx,)))

for p, t in zip(tree_leaves(primals), tree_leaves(tangents)):
    if hasattr(p, 'dtype'):
        assert core.primal_dtype_to_tangent_dtype(p.dtype) == t.dtype, f'tangent dtype mismatch for {p.dtype} -> {t.dtype}'

def tangent_dtype_ok(p, t):
    return not hasattr(p, 'dtype') or jax.core.primal_dtype_to_tangent_dtype(p.dtype) == t.dtype

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• jvp called with different primal and tangent shapes;Got prim
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Custom JVP rule {jvp_name} for function {primal_name} must p

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/7a70ddde6397e6d7. Report an issue: GitHub.