jax-ml/jax · error · ValueError

Accumulator aval mismatch: expected {aval}, got {acc.aval}

Error message

Accumulator aval mismatch: expected {aval}, got {acc.aval}

What it means

Utility check used by VJP accumulation: an accumulator object's aval (shape/dtype of what it accumulates) must be type-compatible with the expected cotangent aval. Raised when a pre-built accumulator was made for a different shape or dtype than the gradient being accumulated.

Source

Thrown at jax/_src/api.py:1835

                   for path, thing1, thing2, explanation
                   in equality_errors_pytreedef(in_tree, refs_tree))
  raise ValueError(msg)

def _vjp_not_saveable_error(jaxpr, in_tree, idxs):
  msg = """the VJP function was applied before restoring its not-saveable residuals.

Because `saveable_args` was passed to `jax.vjp`, some argument values that
would have been saved for the backward pass were instead replaced with
`NotSaveable()` sentinels. Before the VJP function can be applied, these
values must be restored, e.g. by assigning to the VJP function's `args_res`
attribute. The values not yet restored correspond to:
"""
  msg += '\n'.join(f"  * {_vjp_arg_name(jaxpr, in_tree, idx)};" for idx in idxs)
  raise ValueError(msg)

def check_accum(aval, acc):
  if not core.typecompat(acc.aval, aval):
    raise ValueError(f"Accumulator aval mismatch: expected {aval}, got {acc.aval}")
  return acc

def _vjp3_bwd(in_tree, out_tree, out_zeros, jaxpr, out_primal_avals, want_logs,
              residuals, structured_res, maybe_accums, out_ct):
  cts_flat, out_tree_ = tree_flatten(out_ct, is_leaf=lambda x: isinstance(x, ad.Zero))
  if out_tree != out_tree_:
    _vjp_ct_tree_error(jaxpr, out_tree, out_tree_)
  _vjp_check_ct_avals(cts_flat, out_primal_avals)
  cts_flat = [ct for ct, k in zip(cts_flat, out_zeros) if not k]
  primals_in = [*maybe_accums, *tree_leaves(structured_res)]
  logs = ad.backward_pass3(jaxpr, True, residuals, primals_in, cts_flat)
  arg_cts = [x.freeze() if isinstance(x, ad.ValAccum) else
             DidntWant() if isinstance(x, ad.NullAccum) else GradRef()
             for x in maybe_accums]
  arg_cts = map(ad.instantiate_zeros, arg_cts)
  arg_cts = tree_unflatten(in_tree, arg_cts)
  return (arg_cts, logs) if want_logs else arg_cts

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct accumulators from the matching primal via aval.to_ct_aval() / zeros_like of the actual argument
2. Clear cached accumulators whenever parameter shapes/dtypes change
3. Prefer the public with_refs API over manually building accumulators

Example fix

# before
acc = ad.ValueAccum(jnp.zeros(n, jnp.float16))  # param is f32 (n, m)
# after
acc = ad.ValueAccum(jnp.zeros_like(param))

assert core.typecompat(acc.aval, expected_aval), f'accumulator {acc.aval} incompatible with {expected_aval}'

def accum_ok(acc, aval): return core.typecompat(acc.aval, aval)

When it happens

Trigger: Passing an ad.Accum (e.g. ad.ValueAccum/RefAccum) built from a differently-shaped or differently-typed zeros array into _vjp_accum/check_accum; internal API misuse or custom accumulation plumbing.

Common situations: Writing custom autodiff plumbing around jax.ad accumulators; dtype policy mismatches (f32 vs f16 accumulation buffers); shape changes after refactors while accumulators are cached.

Related errors

• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• Acc ref dtype must be float32 or int32, got {dtype}
• WGMMA instruction only supports f32, f16 and s32 out (got {o
• primal and tangent arguments to jax.jvp do not match; dtypes

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