jax-ml/jax · error · TypeError

Custom VJP bwd rule {self.bwd} was registered with defvjp_wi

Error message

Custom VJP bwd rule {self.bwd} was registered with defvjp_with_logs and so must produce a pair (in_cts, logs), but got {in_cts}.

What it means

A custom VJP backward rule registered with hijax's defvjp_with_logs must return a pair (in_cts, logs). The wrapper checked the bwd return value and found it was not a 2-element list/tuple.

Source

Thrown at jax/_src/hijax.py:783

      out_nzs = tree_unflatten(self.out_tree, out_nzs_flat)
      out = tree_unflatten(self.out_tree, out_flat)
      return out, res, out_nzs
    else:
      return out, res

  def vjp_bwd_retval(self, res, out_ct):
    static_args = tuple(x.val for x in self.in_avals if isinstance(x, Static))
    in_avals_ = tuple(x for x in self.in_avals if not isinstance(x, Static))
    leaf = lambda x: isinstance(x, ad_util.Zero)
    if self.symbolic_zeros:
      out_ct = tree_map(ad_util.replace_internal_symbolic_zeros, out_ct, is_leaf=leaf)
    else:
      out_ct = tree_map(ad_util.instantiate, out_ct, is_leaf=leaf)
    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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Return `(in_cts, None)` or `(in_cts, logs_dict)` from the bwd rule
If logging is not needed, register with plain defvjp instead of defvjp_with_logs

Example fix

# before
def bwd(res, ct):
    return cts
# after
def bwd(res, ct):
    return cts, None

Defensive patterns

Strategy: try-catch

Validate before calling

ret = bwd(*static_args, res, out_ct)
assert isinstance(ret, (list, tuple)) and len(ret) == 2, 'defvjp_with_logs bwd must return (in_cts, logs)'

Type guard

def bwd_returns_pair(bwd, static_args, res, ct) -> bool:
    r = bwd(*static_args, res, ct)
    return isinstance(r, (list, tuple)) and len(r) == 2

Try / catch

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'must produce a pair' in str(e):
        raise RuntimeError('bwd registered with logs must return (in_cts, logs)') from e
    raise

Prevention

When using defvjp_with_logs, always return two elements
Prefer plain defvjp unless logs are needed

When it happens

Trigger: Calling grad/vjp over a function whose custom_vjp was registered via defvjp_with_logs, where the bwd function returns only in_cts (or some other shape).

Common situations: Upgrading existing custom_vjp code to defvjp_with_logs but forgetting to add the logs element to the bwd return.

Related errors

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