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 primal and tangent outputs with equal container (pytree) structures, but got {self.out_tree} and {tree} respectively.

What it means

The primal and tangent outputs of a custom JVP rule must have identical pytree structure. JAX flattens the tangent output (treating Zero/SymbolicZero as leaves) and compares its tree to the primal's; any structural difference raises this error showing both trees.

Source

Thrown at jax/_src/hijax.py:1054

    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
    args_flat = tree_leaves_checked(self.in_tree, args)
    is_lin = [isinstance(x, ad.GradAccum) for x in args_flat]
    vals = [x for x, l in zip(args_flat, is_lin) if not l]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return tangent structure mirroring the primal exactly, using ad.InstantiatedZero / jnp.zeros-like values (or ad_util zeros) for non-differentiated outputs
2. Construct both via tree_map over the same template: primal_out, tangent_out = tree_map(...), tree_map(...) on the same structure
3. Check with jax.tree_util.tree_structure(out) == tree_structure(out_tangent) in tests

Example fix

// before
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return (x * 2, x + 1), 2 * xd   # tangent not a 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
out, out_t = f_jvp(primals, tangents)
assert jtu.tree_structure(out) == jtu.tree_structure(out_t)

When it happens

Trigger: A defjvp whose tangent output has different nesting than the primal output — e.g. primal returns (a, b) but tangent is a single array, or tangent omits a None entry present in the primal.

Common situations: Rules that compute tangents only for some outputs and return a partial structure; using jax.tree_util.tree_map on one side but not the other; asymmetry introduced when adding a new output to the function.

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