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 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 {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    {str(ty_tree ).replace("'", "")}\nwhile the custom_jvp-decorated function {primal_name} had output container/pytree structure:\n    {str(ty_tree_).replace("'", "")}.

What it means

The first element of a custom_jvp rule's returned pair must match the container/pytree structure of the original function's output (which JAX knows from a prior primal trace). This error reports the two differing structures side by side.

Source

Thrown at jax/_src/custom_derivatives.py:344

  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 "
           "container/pytree structure:\n"
           f"""    {str(ty_tree_).replace("'", "")}.""")
      raise TypeError(m)
    if not all(map(core.typematch, primal_avals, primal_avals_)):
      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 with leaves of the same shape/dtype), but "
           "instead the JVP rule output's first element had shapes/dtypes of:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_jvp-decorated function {primal_name} had output "
           "shapes/dtypes of:\n"
           f"""    {str(ty_tree_).replace("'", "")}""")
      raise TypeError(m)
  primal_avals_out = [core.typeof(x).strip_weak_type() for x in primals_out]
  expected_tangent_avals_out = [
    core.typeof(x).strip_weak_type().to_tangent_aval()
    for x in primals_out]
  tangent_avals_out = [core.typeof(t).strip_weak_type()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the rule's first returned element structurally identical to what the undecorated function returns
2. Run the function once without grad to inspect its output structure and replicate it

Example fix

# before
return [y], dy
# after
return y, dy

import jax
orig_tree = jax.tree_util.tree_structure(f(*sample_args))
# assert rule's first return matches orig_tree before returning

When it happens

Trigger: The rule returns the primal output wrapped differently than the original function does — e.g. original returns a scalar, rule returns a 1-tuple, or nesting levels differ.

Common situations: Rules written before the primal function was refactored to return a different container; generic rule code that always wraps outputs in lists.

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/399ab5dc08f88cb8. Report an issue: GitHub.