jax-ml/jax · error · ValueError

Preinitialized WGMMAAccumulatorRef only supported in pl.run_

Error message

Preinitialized WGMMAAccumulatorRef only supported in pl.run_state.

What it means

WGMMAAccumulatorRef only supports a preinitialized initial value (init=...) inside pl.run_state; anywhere else, get_ref_aval rejects a non-uninitialized _init because preinitialized accumulator refs aren't representable as plain register arrays.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1556

  def get_ref_aval(self) -> state.AbstractRef:
    ty = ClusterBarrierType(
        collective_axes=self.collective_axes,
        num_arrivals=self.num_arrivals,
        orders_tensor_core=self.orders_tensor_core,
        leader_tracked=self.leader_tracked,
    )
    return state.AbstractRef(jax_core.ShapedArray(self.num_barriers, ty), SMEM)


@dataclasses.dataclass(frozen=True)
class WGMMAAccumulatorRef:
  shape: tuple[int, int]
  dtype: jnp.dtype = jnp.float32
  _init: Any = state_types.uninitialized

  def get_ref_aval(self) -> state.AbstractRef:
    if self._init is not state_types.uninitialized:
      raise ValueError(
          "Preinitialized WGMMAAccumulatorRef only supported in pl.run_state."
      )
    return WGMMAAbstractAccumulatorRef(
        jax_core.ShapedArray(shape=self.shape, dtype=self.dtype), MemorySpace.REGS
    )

  @staticmethod
  def init(array):
    return WGMMAAccumulatorRef(array.shape, array.dtype, _init=array)


def _wgmma_ref_type_mapping(ref: WGMMAAccumulatorRef):
  aval = WGMMAAbstractAccumulatorRef(
      jax_core.ShapedArray(shape=ref.shape, dtype=ref.dtype), MemorySpace.REGS
  )
  return aval, ref._init
state_types._ref_type_aval_mappings[WGMMAAccumulatorRef] = _wgmma_ref_type_mapping

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the kernel in pl.run_state so preinitialized accumulator refs are supported
2. Drop the init argument (leave the accumulator uninitialized) if run_state isn't used
3. Initialize the accumulator by a separate explicit copy/add inside the kernel

Example fix

// before
acc_ref = plgpu.WGMMAAccumulatorRef((m, n), init=x)  # aval outside run_state -> ValueError
// after
@pl.run_state
def kernel(...):
    acc_ref = plgpu.WGMMAAccumulatorRef((m, n), init=x)
    ...

if init is not None:
    assert in_run_state_context, 'WGMMA init= requires pl.run_state'

When it happens

Trigger: Constructing WGMMAAccumulatorRef(..., init=some_value) and then obtaining its ref aval outside pl.run_state — e.g. via a JAX transform tracing the kernel, or building the aval manually.

Common situations: Using init= to seed WGMMA accumulators in Mosaic kernels not wrapped in pl.run_state; tracing/transforming kernels that contain preinitialized accumulators.

Related errors

• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}
• mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}
• The minor dimension size of an accumulator ref must be {info

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