jax-ml/jax · error · TypeError

{type(_prim).__name__} returned structured residuals from `v

Error message

{type(_prim).__name__} returned structured residuals from `vjp_fwd`, which requires overriding `vjp_bwd(res, sres, outgrad, *arg_accums)`

What it means

A HiPrim's vjp_fwd returned structured residuals (a second residual tree beyond flat residuals), which is only supported if the subclass overrides the accumulator-style `vjp_bwd(res, sres, outgrad, *arg_accums)` signature; the default HiPrim.vjp_bwd does not accept sres.

Source

Thrown at jax/_src/hijax.py:426

# A `lin` or `vjp_fwd` rule may return (ans, res), (ans, res, nzs_out), or
# (ans, res, nzs_out, sres). When it returns structured residuals, the paired
# backward rule receives them explicitly: `linearized(res, sres, *tangents)`,
# and `vjp_bwd(res, sres, outgrad, *arg_accums)` (which must be overridden).
def _call_hi_primitive_linearize(is_vjp, nz_in_flat, *args_flat, _prim):
  args = tree_unflatten(_prim.in_tree, args_flat)
  nzs_in = tree_unflatten(_prim.in_tree, nz_in_flat)
  if is_vjp:
    ans, residuals, *rest = _prim.vjp_fwd(nzs_in, *args)
    linearized = partial(fake_linear_op, _prim, nz_in_flat)
  else:
    ans, residuals, *rest = _prim.lin(nzs_in, *args)
    linearized = partial(flatten_user_linearized, _prim)
  ans_flat = tree_leaves_checked(_prim.out_tree, ans)
  nzs_out = rest[0] if rest else True
  sres = rest[1] if len(rest) > 1 else None
  if (sres is not None and is_vjp and
      type(_prim).vjp_bwd is HiPrim.vjp_bwd):
    raise TypeError(
        f"{type(_prim).__name__} returned structured residuals from `vjp_fwd`, "
        "which requires overriding `vjp_bwd(res, sres, outgrad, *arg_accums)`")
  nzs_out_flat = broadcast_prefix(nzs_out, ans)
  linearized = partial(linearized, nzs_out_flat) if is_vjp else linearized
  return ans_flat, nzs_out_flat, residuals, sres, linearized
ad.primitive_linearizations[call_hi_primitive_p] = _call_hi_primitive_linearize

def fake_linear_op(prim, nz_in_flat, nz_out_flat, rs, sres, *tangents):
  rs = rs if sres is None else (rs, sres)  # unpacked in the transpose rule
  residuals_flat, residuals_tree = tree_flatten(rs)
  assert nz_in_flat == [not isinstance(t, ad_util.Zero) for t in tangents]
  nz_tangents = tree_leaves(tangents)
  out_nz = call_hi_primitive_linearized_p.bind(
      *residuals_flat, *nz_tangents, residuals_tree=residuals_tree, _prim=prim,
      nz_in_flat=tuple(nz_in_flat), nz_out_flat=tuple(nz_out_flat),
      has_sres=sres is not None)
  out_nz_iter = iter(out_nz)
  out = [next(out_nz_iter) if nz else ad_util.Zero(a.to_tangent_aval())

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Override `def vjp_bwd(self, res, sres, outgrad, *arg_accums)` on the subclass
2. Or stop returning structured residuals from vjp_fwd (return only flat residuals)
3. Or use vjp_bwd_retval with the matching sres-aware flatten path

Example fix

class MyPrim(hijax.HiPrim):
  def vjp_fwd(self, *args):
    out, res, sres = ...
    return out, res, sres
  # after: add sres-aware backward
  def vjp_bwd(self, res, sres, outgrad, *accums):
    ...

from jax._src.hijax import HiPrim
assert not (returns_sres and type(prim).vjp_bwd is HiPrim.vjp_bwd), \
    'structured residuals require overriding vjp_bwd(res, sres, outgrad, *accums)'

def sres_supported(p) -> bool:
    from jax._src.hijax import HiPrim
    return type(p).vjp_bwd is not HiPrim.vjp_bwd

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'structured residuals' in str(e):
        raise RuntimeError('implement vjp_bwd with sres or drop sres from vjp_fwd') from e
    raise

When it happens

Trigger: A subclass vjp_fwd returns more than two values (out, residuals, sres...) and the class still uses the inherited vjp_bwd, during linearize staging of a vjp.

Common situations: Upgrading a custom primitive to keep structured residual state (e.g. per-layer caches) without updating the backward rule signature.

Related errors

• {type(_prim).__name__}.vjp_bwd should return None or a dict
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• structure of the differentiated function {jaxpr.debug_info.f
• Pure callbacks do not support transpose. Please use `jax.cus

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