{"record":{"id":"aa67f62150186847","repo":"jax-ml/jax","slug":"non-trivial-indexing-on-wgmmaabstractaccumulatorre","errorCode":null,"errorMessage":"Non-trivial indexing on WGMMAAbstractAccumulatorRef is not supported for stores.","messagePattern":"Non-trivial indexing on WGMMAAbstractAccumulatorRef is not supported for stores\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1600,"sourceCode":"  def update(self, inner_aval=None, memory_space=None, kind=None):\n    ref = super().update(inner_aval, memory_space, kind)\n    return WGMMAAbstractAccumulatorRef(\n        inner_aval=ref.inner_aval,\n        memory_space=ref.memory_space,\n    )\n\n  def _getitem(self, tracer, idx):\n    from jax._src.pallas.mosaic_gpu.primitives import wgmma_accumulator_load  # pyrefly: ignore[missing-import]\n    arr = wgmma_accumulator_load(tracer, wait_n=0)\n    if not is_trivial_index(idx, tracer.shape):\n      arr = arr[idx]\n\n    return arr\n\n  def _setitem(self, tracer, idx, value):\n    from jax._src.pallas.mosaic_gpu.primitives import wgmma_accumulator_store  # pyrefly: ignore[missing-import]\n    if not is_trivial_index(idx, tracer.shape):\n      raise NotImplementedError(\n          \"Non-trivial indexing on WGMMAAbstractAccumulatorRef is not supported\"\n          \" for stores.\"\n      )\n    wgmma_accumulator_store(tracer, value)\n\n\nclass AbstractTMEMRef(state.AbstractRef):\n  __slots__ = [\"inner_aval\", \"memory_space\", \"layout\", \"collective\"]\n\n  def __init__(self, inner_aval, memory_space, layout, collective):\n    super().__init__(inner_aval, memory_space)\n    self.layout = layout\n    self.collective = collective\n\n  def __repr__(self) -> str:\n    return f'TMEM({self.inner_aval.str_short()}, layout={self.layout}, collective={self.collective})'\n\n  def update(self, inner_aval=None, memory_space=None, kind=None):","sourceCodeStart":1582,"sourceCodeEnd":1618,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1582-L1618","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nacc_ref[:, 1] = partial_result  # non-trivial index\n\n// after\nfull = block_full_result(acc_ref)  # compute over whole tile\nacc_ref[:] = full  # trivial full-tile store","handlingStrategy":"fallback","validationCode":"from jax._src.pallas.mosaic_gpu.core import is_trivial_index\n# before storing:\nassert is_trivial_index(idx, acc_ref.shape), 'use full-tile stores on WGMMA accumulator'","typeGuard":"def is_full_accumulator_store(idx, shape) -> bool:\n    from jax._src.pallas.mosaic_gpu.core import is_trivial_index\n    return is_trivial_index(idx, shape)","tryCatchPattern":"try:\n    acc_ref[idx] = value\nexcept NotImplementedError:\n    # fall back to computing and storing the full tile\n    acc_ref[:] = full_tile_value","preventionTips":["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"],"tags":["jax","pallas","mosaic-gpu","wgmma","not-implemented","indexing"],"backgroundTag":"unsupported-operation-indexing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}