jax-ml/jax · error · TypeError

Custom JVP rule {jvp_name} for function {primal_name} must p

Error message

Custom JVP rule {jvp_name} for function {primal_name} must produce primal and tangent outputs with equal container (pytree) structures, but got {out_tree} and {out_tree2} respectively.

What it means

The primal and tangent outputs returned by a custom_jvp rule must have identical pytree container structure. This fires when, e.g., the primal is a tuple of two arrays but the tangent is a single array or a differently nested tuple.

Source

Thrown at jax/_src/custom_derivatives.py:323

def _flatten_jvp(f, store, primal_name, jvp_name, in_tree, maybe_out_type, *args):
  primals_in, tangents_in = split_list(args, [len(args) // 2])
  py_primals = tree_unflatten(in_tree, primals_in)
  py_tangents = tree_unflatten(in_tree, tangents_in)
  pair_out = f(py_primals, py_tangents)
  if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:
    msg = (f"Custom JVP rule {jvp_name} for function {primal_name} "
           "must produce a pair (list or tuple of length two) representing "
           f"primal and tangent outputs, but got {pair_out}.")
    raise TypeError(msg)
  py_primals_out, py_tangents_out = pair_out
  primals_out, out_tree = tree_flatten(py_primals_out)
  tangents_out, out_tree2 = tree_flatten(py_tangents_out)
  primal_avals = [core.typeof(x) for x in primals_out]
  if out_tree != out_tree2:
    msg = (f"Custom JVP rule {jvp_name} for function {primal_name} must "
           "produce primal and tangent outputs with equal container (pytree) "
           f"structures, but got {out_tree} and {out_tree2} respectively.")
    raise TypeError(msg)
  # If the primal function already ran, check out_tree agreement.
  try: out_type_ = maybe_out_type()
  except lu.StoreException: out_type_ = None
  if out_type_ is not None:
    out_tree_, primal_avals_, () = out_type_
    ty_tree  = tree_unflatten(out_tree , [a.str_short() for a in primal_avals])
    ty_tree_ = tree_unflatten(out_tree_, [a.str_short() for a in primal_avals_])
    if out_tree_ != out_tree:
      m = (f"Custom JVP rule {jvp_name} for function {primal_name} must "
           "produce a pair (list or tuple of length two) "
           "where the first element represents the primal output "
           "(equal in value to the output of the custom_jvp-decorated function "
           f"{primal_name}, "
           "and in particular of the same container/pytree structure), but "
           "instead the JVP rule output's first element had container/pytree "
           "structure:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_jvp-decorated function {primal_name} had output "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build the tangent as the same pytree as the primal, e.g. tree_map(jnp.zeros_like, primal_out) then fill in real tangents
2. Mirror tuple nesting exactly between the two returned values

Example fix

# before
return (a, b), dz
# after
return (a, b), (dz, jnp.zeros_like(b))

import jax
assert jax.tree_util.tree_structure(primal_out) == jax.tree_util.tree_structure(tangent_out)

When it happens

Trigger: A defjvp rule whose primal output is (a, b) but whose tangent is a scalar zeros; mixing list/tuple nesting between the two halves of the returned pair.

Common situations: Multi-output functions where only one output's tangent was implemented; constructing tangents with jnp.zeros_like applied to the wrong container.

Related errors

• Custom JVP rule {jvp_name} for function {primal_name} must p
• Custom JVP rule must produce primal and tangent outputs with
• primal and tangent arguments to jax.jvp must have the same t
• structure of the differentiated function {jaxpr.debug_info.f
• Pure callbacks do not support JVP. Please use `jax.custom_jv

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