jax-ml/jax · error · TypeError

primal and tangent arguments to jax.jvp must have the same t

Error message

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}.

What it means

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.

Source

Thrown at jax/_src/api.py:1456

  >>>
  >>> primals, tangents = jax.jvp(jax.numpy.sin, (0.1,), (0.2,))
  >>> print(primals)
  0.09983342
  >>> print(tangents)
  0.19900084
  """
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct tangents from primals with jax.tree.map(jnp.ones_like, primals) or jax.lax.zeros_like_pytree so structures always match
2. Fix the container types/keys so both sides match exactly
3. Add asserts comparing tree structures before calling jvp

Example fix

# before
jax.jvp(f, (x, y), (t_x,))
# after
tangents = jax.tree.map(jnp.zeros_like, (x, y))
jax.jvp(f, (x, y), tangents)

assert jax.tree.structure(primals) == jax.tree.structure(tangents), 'primal/tangent tree mismatch'

def trees_match(p, t): return jax.tree.structure(p) == jax.tree.structure(t)

When it happens

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

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

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim
• expected {in_avals.tree}, got {tangents_ft.tree}
• structure of the differentiated function {jaxpr.debug_info.f

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