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) 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.

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)

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 {jvp_name} for function {self.traced.fun_nam
• 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
• The input arguments to the custom_jvp-decorated function {pr

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