jax-ml/jax · error · TypeError

Custom JVP rule for function must produce a pair (list or…

Error message

Custom JVP rule {jvp_name} for function {self.traced.fun_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 {self.traced.fun_name}, and in particular with leaves of the same shape/dtype), but instead the JVP rule output's first element had shapes/dtypes of:
    {str(ty_tree ).replace("'", "")}
while the custom_jvp-decorated function {self.traced.fun_name} had output shapes/dtypes of:
    {str(ty_tree_).replace("'", "")}

What it means

The first element (primal output) of the custom JVP rule's return pair has a different pytree structure than the output of the @jax.custom_jvp-decorated function. JAX flattens both and compares trees; a mismatch (extra nesting, dict vs tuple, None placement) triggers this error with a detailed diff of shapes/dtypes.

Solutions

1. Make the JVP rule's first element structurally identical to the decorated function's return (same tuple nesting, keys, Nones)
2. Wrap/unpack outputs symmetrically: if f returns (a, b), the rule must return (a, b), tan_a, tan_b as (primal_pair, tangent_pair)
3. Add a quick unit test comparing jax.tree.structure(f(x)) with the rule's first output

Example fix

// before
@jax.custom_jvp
def f(x):
  return x * 2, x + 1
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return x * 2, 2 * xd        # primal is single array, should be pair
// after
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return (x * 2, x + 1), (2 * xd, xd)

import jax.tree_util as jtu
struct_f = jtu.tree_structure(jax.eval_shape(f, x))
struct_rule = jtu.tree_structure(jax.eval_shape(lambda p, t: f_jvp(p, t)[0], (x,), (t,)))
assert struct_f == struct_rule

When it happens

Trigger: A defjvp returning e.g. a flat array where the decorated function returns (array,), or returning a dict where the function returns a tuple, or omitting a None output.

Common situations: Decorated function returns multiple outputs but the JVP rule returns one packed object; refactoring output structure without updating the rule; mixing tuple/list conventions inconsistently across the pair.

Related errors

• Custom JVP rule for function must produce primal and…
• Can't use ``defjvps`` with ``nondiff_argnums``.
• Can't use ``defjvps`` with ``nondiff_argnums``.
• custom_jvp-decorated function
• custom_jvp inputs marked with nondiff_argnums must be…

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

Appendix: source

Thrown at jax/_src/hijax.py:1050

    primals_ = tuple(x for x in primals if not isinstance(x, Static))
    tangents_ = tuple(t for x, t in zip(primals, tangents)
                      if not isinstance(x, Static))
    zero = lambda x: isinstance(x, ad_util.Zero)
    if self.symbolic_zeros:
      tangents_ = tree_map(ad_util.replace_internal_symbolic_zeros, tangents_,
                           is_leaf=zero)
    else:
      tangents_ = tree_map(ad_util.instantiate, tangents_, is_leaf=zero)
    pair_out = self.jvp_fun(*static_args, primals_, tangents_)
    jvp_name = getattr(self.jvp_fun, '__name__', str(self.jvp_fun))
    if not isinstance(pair_out, (list, tuple)) or len(pair_out) != 2:
      raise TypeError(
          f"Custom JVP rule {jvp_name} for function {self.traced.fun_name} "
          "must produce a pair (list or tuple of length two) representing "
          f"primal and tangent outputs, but got {pair_out}.")
    out, out_tangent = pair_out
    if (tree := tracing_registry.flatten(out)[1]) != self.out_tree:
      raise TypeError(_jvp_primal_tree_mismatch_err(self, jvp_name, out))
    _jvp_check_primal_avals(self, jvp_name, out)
    zero_ = lambda x: isinstance(x, (ad_util.Zero, ad_util.SymbolicZero))
    if (tree := tracing_registry.flatten(out_tangent, zero_)[1]) != self.out_tree:
      raise TypeError(
          f"Custom JVP rule {jvp_name} for function {self.traced.fun_name} "
          "must produce primal and tangent outputs with equal container "
          f"(pytree) structures, but got {self.out_tree} and {tree} "
          "respectively.")
    _jvp_check_tangent_avals(self, out, out_tangent)
    out_tangent = tree_map(ad_util.replace_rule_output_symbolic_zeros,
                           out_tangent, is_leaf=zero_)
    return out, out_tangent

  lin, linearized = linearize_from_jvp
  vjp_fwd, vjp_bwd_retval = vjp_from_jvp

  def transpose(self, out_ct, *args):
    # The application must be linear in the accumulated args

View on GitHub (pinned to 1e1c6a8fc0)