{"record":{"id":"47cb81d568356b41","repo":"jax-ml/jax","slug":"accumulator-aval-mismatch-expected-aval-got-a","errorCode":null,"errorMessage":"Accumulator aval mismatch: expected {aval}, got {acc.aval}","messagePattern":"Accumulator aval mismatch: expected (.+?), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1835,"sourceCode":"                   for path, thing1, thing2, explanation\n                   in equality_errors_pytreedef(in_tree, refs_tree))\n  raise ValueError(msg)\n\ndef _vjp_not_saveable_error(jaxpr, in_tree, idxs):\n  msg = \"\"\"the VJP function was applied before restoring its not-saveable residuals.\n\nBecause `saveable_args` was passed to `jax.vjp`, some argument values that\nwould have been saved for the backward pass were instead replaced with\n`NotSaveable()` sentinels. Before the VJP function can be applied, these\nvalues must be restored, e.g. by assigning to the VJP function's `args_res`\nattribute. The values not yet restored correspond to:\n\"\"\"\n  msg += '\\n'.join(f\"  * {_vjp_arg_name(jaxpr, in_tree, idx)};\" for idx in idxs)\n  raise ValueError(msg)\n\ndef check_accum(aval, acc):\n  if not core.typecompat(acc.aval, aval):\n    raise ValueError(f\"Accumulator aval mismatch: expected {aval}, got {acc.aval}\")\n  return acc\n\ndef _vjp3_bwd(in_tree, out_tree, out_zeros, jaxpr, out_primal_avals, want_logs,\n              residuals, structured_res, maybe_accums, out_ct):\n  cts_flat, out_tree_ = tree_flatten(out_ct, is_leaf=lambda x: isinstance(x, ad.Zero))\n  if out_tree != out_tree_:\n    _vjp_ct_tree_error(jaxpr, out_tree, out_tree_)\n  _vjp_check_ct_avals(cts_flat, out_primal_avals)\n  cts_flat = [ct for ct, k in zip(cts_flat, out_zeros) if not k]\n  primals_in = [*maybe_accums, *tree_leaves(structured_res)]\n  logs = ad.backward_pass3(jaxpr, True, residuals, primals_in, cts_flat)\n  arg_cts = [x.freeze() if isinstance(x, ad.ValAccum) else\n             DidntWant() if isinstance(x, ad.NullAccum) else GradRef()\n             for x in maybe_accums]\n  arg_cts = map(ad.instantiate_zeros, arg_cts)\n  arg_cts = tree_unflatten(in_tree, arg_cts)\n  return (arg_cts, logs) if want_logs else arg_cts\n","sourceCodeStart":1817,"sourceCodeEnd":1853,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1817-L1853","documentation":"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.","triggerScenarios":"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.","commonSituations":"Writing custom autodiff plumbing around jax.ad accumulators; dtype policy mismatches (f32 vs f16 accumulation buffers); shape changes after refactors while accumulators are cached.","solutions":["Construct accumulators from the matching primal via aval.to_ct_aval() / zeros_like of the actual argument","Clear cached accumulators whenever parameter shapes/dtypes change","Prefer the public with_refs API over manually building accumulators"],"exampleFix":"# before\nacc = ad.ValueAccum(jnp.zeros(n, jnp.float16))  # param is f32 (n, m)\n# after\nacc = ad.ValueAccum(jnp.zeros_like(param))","handlingStrategy":"type-guard","validationCode":"assert core.typecompat(acc.aval, expected_aval), f'accumulator {acc.aval} incompatible with {expected_aval}'","typeGuard":"def accum_ok(acc, aval): return core.typecompat(acc.aval, aval)","tryCatchPattern":null,"preventionTips":["Build accumulators from the corresponding primal's to_ct_aval","Invalidate accumulator caches on shape/dtype changes","Prefer public with_refs over manual ad.Accum construction"],"tags":["jax","vjp","accumulator","internal-api","dtype"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}