jax-ml/jax · error · ValueError

{result}the bwd rule attached to {primal_sourceinfo} produce

Error message

{result}the bwd rule attached to {primal_sourceinfo} produced an output of type {ct_aval.str_short()} which doesn't match expected type {expected.str_short()}

What it means

The backward (VJP) rule attached to a @jax.custom_vjp function via defvjp returned cotangents whose avals don't match the expected cotangent avals of the primal inputs. JAX validates each bwd output type against primal_aval.to_ct_aval() (skippable by setting JAX_DISABLE_BWD_CHECKS or the config disable_bwd_checks).

Source

Thrown at jax/_src/hijax.py:906

def _vjp_fwd_aval_mismatch_err(path, primal_aval, fwd_val):
  if not core.typematch(ty := typeof(fwd_val), primal_aval):
    raise TypeError(f"at {keystr(path)}, got fwd output type {ty.str_short()} "
                    f"which doesn't match primal output type {primal_aval.str_short()}")

def _vjp_bwd_aval_mismatch_err(primal_sourceinfo, path, primal_aval, ct):
  if config.disable_bwd_checks.value:
    return
  if isinstance(ct, ad_util.Zero):
    return
  if isinstance(primal_aval, AbstractRef):
    primal_aval = primal_aval.inner_aval
  expected = primal_aval.to_ct_aval()
  ct_aval = ct.aval if isinstance(ct, ad_util.SymbolicZero) else typeof(ct)
  if (not core.typematch(expected, ct_aval) and
      not _temporary_dtype_exception(expected, ct_aval) and
      getattr(expected, 'dtype', None) is not dtypes.float0):
    result = f"at output{keystr(path)} " if path else ""
    raise ValueError(
        f"{result}the bwd rule attached to {primal_sourceinfo} produced an"
        f" output of type {ct_aval.str_short()} which doesn't match expected"
        f" type {expected.str_short()}")

def _replace_none(primal_in_aval, maybe_ct):
  if maybe_ct is None:
    return ad_util.Zero(primal_in_aval.to_ct_aval())
  else:
    return maybe_ct

class custom_vjp3:
  fwd: Callable | None = None
  bwd: Callable | None = None
  symz: bool = False
  opt_remat: bool = False
  with_logs: bool = False

  def __init__(self, f, nondiff_argnums=(), nondiff_argnames=()):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make bwd return cotangents matching each diff-argnum input's shape and dtype (use jnp.zeros(x.shape, x.dtype) pattern)
Cast bwd outputs with jax.lax.convert_element_type to the input dtype
As a last-resort diagnostic escape hatch, set jax.config.update('jax_disable_bwd_checks', True) (does not fix the underlying wrong gradients)
Verify pytree structure of bwd's return matches the diff-argnums prefix of f's inputs

Example fix

// before
def f_bwd(res, ct):
  return (jnp.zeros_like(ct.astype(jnp.float64)),)  # wrong dtype
// after
def f_bwd(res, ct):
  return (ct * res,)  # same shape/dtype as the input x

Defensive patterns

Strategy: validation

Validate before calling

# check bwd cotangent avals against inputs
expected = jax.eval_shape(lambda *args: args[:-0] if False else args, *diff_args)
ct_out = jax.eval_shape(lambda res, ct: f_bwd(res, ct), residuals, dummy_ct)
# structures/dtypes must match diff args

Try / catch

try:
    jax.grad(f)(x)
except ValueError as e:
    if 'bwd rule' in str(e) and "doesn't match expected type" in str(e):
        # fix bwd return shape/dtype to match inputs
        ...

Prevention

Mirror each bwd output to its input's shape/dtype with zeros_like patterns
Add eval_shape-based tests for fwd/bwd rules
Avoid casting inside bwd unless you cast back

When it happens

Trigger: A bwd function returning arrays with wrong shape, dtype (e.g. float64 vs float32), or wrong pytree structure for the input cotangents; returning zeros of the wrong shape; integer float0 handling mistakes.

Common situations: Hand-written bwd rules that return jnp.zeros_like(residual) instead of zeros matching the input; enabling/disabling 64-bit mode after writing the rule; returning None where a cotangent is expected or vice versa.

Related errors

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