jax-ml/jax · error · ValueError

Accumulator dtype {inner.dtype} does not match value dtype {

Error message

Accumulator dtype {inner.dtype} does not match value dtype {val.dtype}

What it means

The WGMMA accumulator store requires the stored value's dtype to exactly match the accumulator's inner dtype (typically fp32). Storing a half/bfloat16 result directly fails.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:2355


@wgmma_accumulator_store_p.def_effectful_abstract_eval
def _wgmma_accumulator_store_abstract_eval(acc, val):
  # Before discharge acc is a WGMMAAbstractAccumulatorRef. After discharge,
  # the discharge rule re-binds the primitive and acc becomes a ShapedArray.
  if isinstance(acc, gpu_core.WGMMAAbstractAccumulatorRef):
    inner = acc.inner_aval
    assert isinstance(inner, jax_core.ShapedArray)
  elif isinstance(acc, jax_core.ShapedArray):
    inner = acc
  else:
    raise TypeError(f"Expected WGMMAAbstractAccumulatorRef or ShapedArray, got {type(acc)}")
  if inner.shape != val.shape:
    raise ValueError(
        f"Accumulator shape {inner.shape} does not match value shape {val.shape}"
    )
  if inner.dtype != val.dtype:
    raise ValueError(
        f"Accumulator dtype {inner.dtype} does not match value dtype {val.dtype}"
    )
  effects: set[jax_core.Effect] = {gpu_core._wgmma_pipeline_effect}
  if isinstance(acc, gpu_core.WGMMAAbstractAccumulatorRef):
    effects.add(state.WriteEffect(0))
  return inner, effects


@discharge.register_discharge_rule(wgmma_accumulator_store_p)
def _wgmma_accumulator_store_discharge(ctx, acc, val):
  del ctx
  return (wgmma_accumulator_store_p.bind(acc, val), None), []


@lowering.register_lowering_rule(
    wgmma_accumulator_store_p, mgpu.LoweringSemantics.Lane
)
def _wgmma_accumulator_store_lowering(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast the value to the accumulator dtype before storing: val.astype(acc.inner_aval.dtype)
Keep all arithmetic on the accumulator value in fp32
If intentional, make the cast explicit rather than relying on implicit promotion

Example fix

# before
wgmma_accumulator_store(acc, (delta).astype(jnp.bfloat16))
# after
wgmma_accumulator_store(acc, delta.astype(acc.inner_aval.dtype))

Defensive patterns

Strategy: validation

Validate before calling

val = val.astype(acc.inner_aval.dtype)
wgmma_accumulator_store(acc, val)

Type guard

def dtype_ok(acc, val):
    inner = getattr(acc, 'inner_aval', acc)
    return inner.dtype == val.dtype

Prevention

Standardize on fp32 for accumulator arithmetic
Cast explicitly at store boundaries

When it happens

Trigger: Computing the update in f16/bf16 (e.g. output of an op that downcasts) and storing into an fp32 WGMMA accumulator, or converting the loaded value with .astype(jnp.float16) before storing back.

Common situations: Mixed-precision kernels where the matmul operands are fp8/fp16 but accumulators are fp32; forgetting to cast the delta back after arithmetic in lower precision.

Related errors

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