jax-ml/jax · error · TypeError

{type(_prim).__name__}.vjp_bwd should return None or a dict

Error message

{type(_prim).__name__}.vjp_bwd should return None or a dict of backward-pass log entries, got {type(log).__name__}

What it means

The transpose of the linearized HiPrim primitive ran vjp_bwd, whose return value must be None or a dict of backward-pass log entries; anything else is a TypeError. The actual cotangents flow through accumulators, so the return is only for logging.

Source

Thrown at jax/_src/hijax.py:492

  accums_flat_ = iter(accums_flat)
  accums_flat = [next(accums_flat_) if nz else ad.NullAccum(aval.to_ct_aval())
                 for aval, nz in zip(_prim.in_avals_flat, nz_in_flat)]
  assert next(accums_flat_, None) is None
  accums = tree_unflatten(_prim.in_tree, accums_flat)
  cts_flat_iter = iter(cts_flat_)
  cts_flat = [next(cts_flat_iter) if nz else ad_util.Zero(a.to_ct_aval())
              for a, nz in zip(_prim.out_avals_flat, nz_out_flat)]
  assert next(cts_flat_iter, sentinel := object()) is sentinel
  cts = tree_unflatten(_prim.out_tree, cts_flat)
  # A vjp_bwd rule may return a dict of pytrees to log out of the backward
  # pass (see VJP.with_logs), or None (the usual case) to log nothing.
  if has_sres:
    residuals, sres = residuals
    log = _prim.vjp_bwd(residuals, sres, cts, *accums)
  else:
    log = _prim.vjp_bwd(residuals, cts, *accums)
  if log is not None and type(log) is not dict:
    raise TypeError(
        f"{type(_prim).__name__}.vjp_bwd should return None or a dict of "
        f"backward-pass log entries, got {type(log).__name__}")
  return log
ad.fancy_transposes[call_hi_primitive_linearized_p] = _call_hi_primitive_linearized_transpose

def _call_hi_primitive_linearized_prettyprint(eqn, context, settings):
  params = dict(eqn.params, _prim=eqn.params['_prim'].__class__.__name__,
                residuals_tree='...')
  if not params['has_sres']:
    del params['has_sres']
  return core._pp_eqn(eqn.replace(params=params), context, settings)
core.pp_eqn_rules[call_hi_primitive_linearized_p] = _call_hi_primitive_linearized_prettyprint

def _call_hi_primitive_jvp(primals, tangents, *, _prim):
  primals = tree_unflatten(_prim.in_tree, primals)
  tangents = tree_unflatten(_prim.in_tree, tangents)
  out_primals, out_tangents = _prim.jvp(primals, tangents)
  out_primals_flat = tree_leaves_checked(_prim.out_tree, out_primals)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change vjp_bwd to return None (or a dict of logs) and deliver cotangents via the arg_accums accumulators
2. If you want the classic return-based API, implement vjp_bwd_retval instead
3. Return logs only as {key: value} entries for backward-pass logging

Example fix

# before
def vjp_bwd(self, res, outgrad, *accums):
    return cts  # wrong
# after
def vjp_bwd(self, res, outgrad, *accums):
    ...
    return None  # or {'loss_scale': ls}

def vjp_bwd_returns_valid(p) -> bool:
    out = p.vjp_bwd(dummy_res, dummy_ct, *dummy_accums)
    return out is None or type(out) is dict

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'vjp_bwd should return' in str(e):
        raise RuntimeError('vjp_bwd must return None or a logs dict; cots go via accumulators') from e
    raise

When it happens

Trigger: A custom vjp_bwd implementation returns the input cotangents (a tuple/list/array) instead of None or a logs dict.

Common situations: Porting a classic custom_vjp bwd rule that returns in_cts into the hijax accumulator API without dropping the return value.

Related errors

• {type(_prim).__name__} returned structured residuals from `v
• 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/b6923a03d80f1179. Report an issue: GitHub.