jax-ml/jax · error · NotImplementedError

Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

Error message

Non-trivial indexing on WGMMAAbstractAccumulatorRef is not supported for stores.

What it means

Raised when a store (assignment) into a WGMMA accumulator Ref uses non-trivial indexing (anything beyond full-dimension slices). The Mosaic GPU WGMMA/TMEM accumulator can only be written back with wgmma_accumulator_store over the whole accumulator, so partial slices, integer indices, or stepped slices are rejected.

Source

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

  def update(self, inner_aval=None, memory_space=None, kind=None):
    ref = super().update(inner_aval, memory_space, kind)
    return WGMMAAbstractAccumulatorRef(
        inner_aval=ref.inner_aval,
        memory_space=ref.memory_space,
    )

  def _getitem(self, tracer, idx):
    from jax._src.pallas.mosaic_gpu.primitives import wgmma_accumulator_load  # pyrefly: ignore[missing-import]
    arr = wgmma_accumulator_load(tracer, wait_n=0)
    if not is_trivial_index(idx, tracer.shape):
      arr = arr[idx]

    return arr

  def _setitem(self, tracer, idx, value):
    from jax._src.pallas.mosaic_gpu.primitives import wgmma_accumulator_store  # pyrefly: ignore[missing-import]
    if not is_trivial_index(idx, tracer.shape):
      raise NotImplementedError(
          "Non-trivial indexing on WGMMAAbstractAccumulatorRef is not supported"
          " for stores."
      )
    wgmma_accumulator_store(tracer, value)


class AbstractTMEMRef(state.AbstractRef):
  __slots__ = ["inner_aval", "memory_space", "layout", "collective"]

  def __init__(self, inner_aval, memory_space, layout, collective):
    super().__init__(inner_aval, memory_space)
    self.layout = layout
    self.collective = collective

  def __repr__(self) -> str:
    return f'TMEM({self.inner_aval.str_short()}, layout={self.layout}, collective={self.collective})'

  def update(self, inner_aval=None, memory_space=None, kind=None):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Store to the full accumulator: compute the complete result and assign `acc_ref[:] = value` (all-dims sliced) or use wgmma_accumulator_store directly
Materialize the accumulator into a regular array (e.g. via convert/load), do arbitrary indexing there, then store back the full result
Restructure the kernel so partial writes go to a normal Ref (SMEM/GMEM) instead of the WGMMA accumulator

Example fix

// before
acc_ref[:, 1] = partial_result  # non-trivial index

// after
full = block_full_result(acc_ref)  # compute over whole tile
acc_ref[:] = full  # trivial full-tile store

Defensive patterns

Strategy: fallback

Validate before calling

from jax._src.pallas.mosaic_gpu.core import is_trivial_index
# before storing:
assert is_trivial_index(idx, acc_ref.shape), 'use full-tile stores on WGMMA accumulator'

Type guard

def is_full_accumulator_store(idx, shape) -> bool:
    from jax._src.pallas.mosaic_gpu.core import is_trivial_index
    return is_trivial_index(idx, shape)

Try / catch

try:
    acc_ref[idx] = value
except NotImplementedError:
    # fall back to computing and storing the full tile
    acc_ref[:] = full_tile_value

Prevention

Only ever assign full-tile results to WGMMA accumulator refs
Keep partial updates on SMEM/GMEM refs, not TMEM accumulators
Wrap accumulator stores in a helper that asserts trivial indexing

When it happens

Trigger: Calling `acc_ref[...] = value` (or any __setitem__) inside a Pallas Mosaic GPU kernel where idx is not a trivial full slice, e.g. `acc_ref[0] = x` or `acc_ref[:, ::2] = y` on a WGMMAAbstractAccumulatorRef obtained from the wgmma accumulate API.

Common situations: Porting a TPU Pallas matmul kernel to GPU WGMMA and reusing slice-assignment code; trying to zero or partially update only part of the accumulator between MMA steps; masking outputs by assigning to a sub-block of the accumulator.

Related errors

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