jax-ml/jax · error · TypeError

Expected WGMMAAbstractAccumulatorRef got {acc_aval}

Error message

Expected WGMMAAbstractAccumulatorRef got {acc_aval}

What it means

The wgmma_ref primitive (updating an accumulator in place) requires its first operand to be a WGMMAAbstractAccumulatorRef — the special accumulator SMEM buffer allocated via the dedicated allocator. Passing any other ref aval is a TypeError.

Source

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

  wgmma_ref_p.bind(
      acc,
      a,
      b,
      *acc_transforms_leaves,
      *a_transforms_leaves,
      *b_transforms_leaves,
      acc_transforms_tree=acc_transforms_tree,
      a_transforms_tree=a_transforms_tree,
      b_transforms_tree=b_transforms_tree,
  )


@wgmma_ref_p.def_effectful_abstract_eval
def _wgmma_ref_effectful_abstract_eval(acc_aval, a_aval, b_aval, *_, **params):
  del b_aval, params
  if not isinstance(acc_aval, gpu_core.WGMMAAbstractAccumulatorRef):
    raise TypeError(f"Expected WGMMAAbstractAccumulatorRef got {acc_aval}")
  return (), {
      gpu_core._wgmma_pipeline_effect,
      state.WriteEffect(0),
      state.ReadEffect(0),
      state.ReadEffect(2),
      *([state.ReadEffect(1)] if isinstance(a_aval, state.AbstractRef) else [])
  }


def _wgmma_ref_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,
):
  del settings
  acc, a, b, *leaves = eqn.invars
  transform_treedefs = [
      eqn.params["acc_transforms_tree"],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the accumulator using the WGMMA accumulator allocation API (mgpu.SMEM_ALLOCATOR / wgmma accumulator helper) so its aval is WGMMAAbstractAccumulatorRef.
2. Check that the accumulator ref is passed as the first positional argument to wgmma.

from jax._src.pallas.mosaic_gpu import gpu_core
def is_wgmma_acc(ref):
    return isinstance(getattr(ref, 'aval', ref), gpu_core.WGMMAAbstractAccumulatorRef)

When it happens

Trigger: Binding wgmma_ref (or the wgmma_increments style API) with an accumulator that is a normal SMEM/TrackedRef instead of a WGMMA accumulator ref.

Common situations: Allocating the accumulator with the generic SMEM allocator instead of the WGMMA accumulator allocator; refactoring kernels where acc refs were swapped with ordinary buffers.

Related errors

• ref must be a reference
• Unsupported type: {x}
• Preinitialized WGMMAAccumulatorRef only supported in pl.run_
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Can only store to references (got {x_ref}).

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