jax-ml/jax · error · ValueError

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

Error message

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

What it means

Raised by the abstract eval of wgmma_accumulator_store when the value being stored into a WGMMA accumulator ref has a different shape than the accumulator's inner ShapedArray. The store must write back a tensor of exactly the same shape that was read from the accumulator.

Source

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

def wgmma_accumulator_store(acc_ref, val):
  if not isinstance(acc_ref.aval, gpu_core.WGMMAAbstractAccumulatorRef):
    raise TypeError(f"acc must be a WGMMAAccumulatorAbstractRef, got {acc_ref.aval=}")
  wgmma_accumulator_store_p.bind(acc_ref, val)


@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), []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Ensure val keeps the exact shape returned by wgmma_accumulator_load
Check any reshape/broadcast between accumulator load and store and remove or compensate for it
Print acc.inner_aval.shape and val.shape right before the store to find the divergence

Example fix

# before
acc_val = wgmma_accumulator_load(acc)
out = acc_val.reshape(other_shape)
wgmma_accumulator_store(acc, out)
# after
acc_val = wgmma_accumulator_load(acc)
out = acc_val + delta  # same shape
wgmma_accumulator_store(acc, out)

Defensive patterns

Strategy: validation

Validate before calling

assert val.shape == acc.inner_aval.shape, (val.shape, acc.inner_aval.shape)
wgmma_accumulator_store(acc, val)

Type guard

def can_store(acc, val):
    inner = getattr(acc, 'inner_aval', acc)
    return getattr(inner, 'shape', None) == val.shape

Try / catch

try:
    wgmma_accumulator_store(acc, val)
except ValueError as e:
    if 'does not match value shape' in str(e):
        wgmma_accumulator_store(acc, val.reshape(inner.shape))
    else:
        raise

Prevention

Keep loaded accumulator values unreshaped until after the store
Add shape asserts right after wgmma_accumulator_load

When it happens

Trigger: Calling wgmma_accumulator_store(acc, val) where val was reshaped/broadcast after the corresponding wgmma_accumulator_load, e.g. acc has shape (128, 8) but the computed value is (64, 16) or has extra leading dims.

Common situations: Reshaping the accumulator value between load and store, using a different BlockMapping/grid that changes block shapes, or passing the result of a matmul whose N dimension differs from the accumulator's.

Related errors

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