jax-ml/jax · error · TypeError

Custom VJP bwd rule {self.bwd} must produce a tuple but got

Error message

Custom VJP bwd rule {self.bwd} must produce a tuple but got {type(in_cts)}.

What it means

After optional log handling, the custom VJP backward rule's in_cts must be a tuple (lists are coerced); any other type raises TypeError before cotangents are aligned to inputs.

Source

Thrown at jax/_src/hijax.py:797

    in_cts = self.bwd(*static_args, res, out_ct)
    logs = None
    if self.with_logs:
      if not (isinstance(in_cts, (list, tuple)) and len(in_cts) == 2):
        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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Return a tuple with one cotangent per non-static primal argument, e.g. `(ct,)` for one argument
Wrap list returns are auto-converted, but prefer tuples explicitly

Example fix

# before
def bwd(res, ct):
    return ct * w
# after
def bwd(res, ct):
    return (ct * w,)

Defensive patterns

Strategy: type-guard

Validate before calling

cts = bwd(*static_args, res, out_ct)
cts = cts if type(cts) is tuple else (cts if isinstance(cts, tuple) else None)
assert isinstance(cts, tuple), 'bwd must return a tuple of cotangents'

Type guard

def cts_is_tuple(bwd_out) -> bool:
    return isinstance(bwd_out, tuple) or isinstance(bwd_out, list)

Try / catch

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'must produce a tuple' in str(e):
        raise RuntimeError('wrap single cotangent: return (ct,)') from e
    raise

Prevention

Always return a tuple, even for one argument
List returns work but tuples are idiomatic

When it happens

Trigger: A bwd rule returning an array, generator, or single value instead of a tuple of per-argument cotangents.

Common situations: bwd written for a single-argument function returns a bare array instead of `(ct,)`.

Related errors

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