jax-ml/jax · error · TypeError

Custom JVP rule {jvp_name} for function {self.traced.fun_nam

Error message

Custom JVP rule {jvp_name} for function {self.traced.fun_name} must produce a pair (list or tuple of length two) representing primal and tangent outputs, but got {pair_out}.

What it means

A custom JVP rule (the function passed to @jax.custom_jvp's defjvp or the jvpfun) must return exactly a length-2 sequence: (primal_out, tangent_out). If it returns anything else (a single array, a 3-tuple, a dict), this TypeError is raised.

Source

Thrown at jax/_src/hijax.py:1044

  def expand(self, *args):
    args = [x for x in args if not isinstance(x, Static)]
    return self.traced(*args)

  def jvp(self, primals, tangents):
    static_args = tuple(x.val for x in primals if isinstance(x, Static))
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return exactly `return primal_out, tangent_out` from the JVP rule
2. Ensure both elements are pytrees matching the decorated function's output structure
3. If tangents are unknown, return zeros of the right shapes (only if semantically correct)

Example fix

// before
@jax.custom_jvp
def f(x):
  return x * 2
@f.defjvp
def f_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return x * 2          # missing tangent
// after
@f.defjvp
def f_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return x * 2, 2. * x_dot

out = f_jvp(*static_args, primals_, tangents_)
assert isinstance(out, (list, tuple)) and len(out) == 2, 'JVP rule must return (primal, tangent)'

def is_rule_pair(v) -> bool:
    return isinstance(v, (list, tuple)) and len(v) == 2

When it happens

Trigger: Writing a defjvp that returns only the primal output, returns a tuple of 3, returns None, or forgets to also compute and return tangents (e.g. returns out, out_tan, extra_debug).

Common situations: First-time custom_jvp users modeling on custom_vjp examples (which have different conventions); refactoring a rule and dropping the tangent return; returning a generator or other non-sequence.

Related errors

• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Can't use ``defjvps`` with ``nondiff_argnums``.
• No JVP defined for custom_jvp function {primal_name} using d
• Custom JVP rule {jvp_name} for function {primal_name} must p
• Custom JVP rule {jvp_name} for function {primal_name} must p

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