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 the second element of the pair it returns must be None or a dict of backward-pass log entries, but got {type(logs).__name__}.

What it means

For a defvjp_with_logs-registered backward rule, the second element of the returned pair must be None or a dict of log entries; another type fails this TypeError after the pair was unpacked.

Source

Thrown at jax/_src/hijax.py:789

  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 "
                       "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:])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Return None or a plain dict, e.g. `(cts, {'grad_norm': g})`
Convert list-based logs to a dict with keys

Example fix

# before
return cts, [g1, g2]
# after
return cts, {'g1': g1, 'g2': g2}

Defensive patterns

Strategy: type-guard

Validate before calling

cts, logs = bwd(*static_args, res, out_ct)
assert logs is None or type(logs) is dict

Type guard

def logs_valid(logs) -> bool:
    return logs is None or type(logs) is dict

Try / catch

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'dict of backward-pass log entries' in str(e):
        raise RuntimeError('second element must be None or dict') from e
    raise

Prevention

Use dicts with named keys for logs
Never return lists/arrays as the logs element

When it happens

Trigger: bwd returns a pair whose second element is a list, string, or array rather than None/dict.

Common situations: Returning a list of log tensors or a scalar metric instead of a named dict of logs.

Related errors

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