jax-ml/jax · error · ValueError

Custom VJP bwd rule {self.bwd} must produce a tuple of lengt

Error message

Custom VJP bwd rule {self.bwd} must produce a tuple of length equal to the primal args tuple, but got length {len(in_cts)}

What it means

The custom VJP backward rule returned a tuple whose length does not match the number of primal (non-static) arguments after two leading None entries for consts/fwd_consts are inserted. ValueError raised during length validation.

Source

Thrown at jax/_src/hijax.py:801

        raise TypeError(
            f"Custom VJP bwd rule {self.bwd} was registered with "
            "defvjp_with_logs and so must produce a pair (in_cts, logs), "
            f"but got {in_cts}.")
      in_cts, logs = in_cts
      if logs is not None and type(logs) is not dict:
        raise TypeError(
            f"Custom VJP bwd rule {self.bwd} was registered with "
            "defvjp_with_logs, and so the second element of the pair it "
            "returns must be None or a dict of backward-pass log entries, "
            f"but got {type(logs).__name__}.")
    if isinstance(in_cts, list):
      in_cts = tuple(in_cts)
    if not isinstance(in_cts, tuple):
      raise TypeError(f"Custom VJP bwd rule {self.bwd} must produce a tuple "
                      f"but got {type(in_cts)}.")
    in_cts = (None, None, *in_cts)  # zero cts for the consts and fwd_consts args
    if len(in_cts) != len(self.in_tree.children()) - len(self.static_argnums):
      raise ValueError(f"Custom VJP bwd rule {self.bwd} must produce a tuple "
                       "of length equal to the primal args tuple, but got "
                       f"length {len(in_cts)}")
    in_cts = broadcast_prefix(in_cts, in_avals_, is_leaf=lambda x: x is None)
    in_cts = tree_unflatten(self.in_tree, map(_replace_none, self.in_avals_flat, in_cts))
    tree_map_with_path(partial(_vjp_bwd_aval_mismatch_err, self.traced._fun_sourceinfo),
                               self.in_avals[2:], in_cts[2:])
    if self.symbolic_zeros:
      in_cts = tree_map(ad_util.replace_rule_output_symbolic_zeros, in_cts)
    return (in_cts, logs) if self.with_logs else in_cts

  def jvp(self, primals, tangents):
    if self.symbolic_zeros: ad.raise_custom_vjp_error_on_jvp()
    zero = lambda x: isinstance(x, ad_util.Zero)
    nzs_in = tuple(tree_map(lambda t: not isinstance(t, ad_util.Zero), t,
                            is_leaf=zero) for t in tangents)
    tangents = tree_map(ad_util.instantiate, tangents, is_leaf=zero)
    if self.opt_remat:
      fwd_traced = api.jit(partial(self.vjp_fwd, nzs_in)).trace(*primals)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Return exactly one cotangent (possibly None) per non-static primal argument
Re-count args after signature changes, including defaults and static args
Use jax.tree.map over the input tree to build the return

Example fix

# before (3 dynamic args)
def bwd(res, ct):
    return (g1,)
# after
def bwd(res, ct):
    return (g1, g2, g3)

Defensive patterns

Strategy: validation

Validate before calling

n_expected = len(prim.in_tree.children()) - len(prim.static_argnums) - 2
cts = bwd(*static_args, res, out_ct)[0] if with_logs else bwd(...)
assert len(cts) == n_expected, f'expected {n_expected} cotangents, got {len(cts)}'

Type guard

def cts_arity_ok(cts, n_dynamic_args) -> bool:
    return len(cts) == n_dynamic_args

Try / catch

try:
    jax.grad(f)(x)
except ValueError as e:
    if 'length equal to the primal args tuple' in str(e):
        raise RuntimeError('bwd must return one ct per non-static primal arg') from e
    raise

Prevention

Update bwd whenever the primal signature changes
Track static_argnums carefully; they reduce the expected ct count
Return None placeholders for args that receive no cotangent

When it happens

Trigger: bwd returns more or fewer cotangents than the function has non-static arguments (e.g. two cts for a three-arg function, or forgetting that some args are static).

Common situations: Adding/removing function arguments after writing the bwd rule; static_argnums in the custom_vjp registration shrinking the expected count.

Related errors

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