jax-ml/jax · error · TypeError

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

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 pairs each primal leaf with a tangent leaf, so both argument trees must have identical structure. Raised after tree-flattening when the pytree structures diverge, e.g. a dict on one side and a tuple on the other, or different numbers of leaves.

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

• Custom JVP rule for function must produce a pair (list or…
• Custom JVP rule for function must produce primal and…
• Custom JVP rule must produce primal and tangent outputs…
• expected , got
• interpret with dynamic grid bounds unsupported

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

Appendix: 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)