jax-ml/jax · error · NotImplementedError

Sparse meta layout stores unsupported.

Error message

Sparse meta layout stores unsupported.

What it means

TensorMem.store refuses to store into TMEM allocated with the sparse metadata layout. The sparse_meta_layout does not faithfully describe the physical TMEM layout written by async_copy_sparse_smem_to_tmem (multiple vector dims would be needed), so a reg->TMEM store could produce inconsistent data.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1479

      raise ValueError(
          f"Stored array has shape {value.shape}, but TMEM has shape"
          f" {self.shape}"
      )
    if value.mlir_dtype != self.dtype:
      raise ValueError(
          f"Stored array has dtype {value.mlir_dtype}, but TMEM has dtype"
          f" {self.dtype}"
      )
    if not isinstance(value.layout, fa.TiledLayout):
      raise TypeError(f"Stored array has layout {value.layout}, but TMEM stores expect a TiledLayout")
    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the
    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.
    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to
    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.
    # This restriction can be lifted if `TiledLayout` supports multiple
    # vector dims.
    if self.layout == sparse_meta_layout():
      raise NotImplementedError("Sparse meta layout stores unsupported.")
    packing = self.packing
    has_default_layout = self.layout == tmem_default_layout(packing=packing)
    bitwidth = utils.bitwidth(self.dtype)
    is_at_least_16b = bitwidth in {16, 32}
    if value.layout == LAYOUT and has_default_layout and is_at_least_16b:
      _store_32xcols(
          self.address, value.registers.T.reshape((4, -1)), packing
      )
    elif value.layout == self.layout.as_tiled_layout() and packing * bitwidth == 32:
      _store_32xcols_native(self.address, value.registers.reshape(-1), packing)
    # TODO(apaszke): Support the case where we have a long vector length in the
    # FA more generally, not just for 2x32b.
    # TODO(apaszke): Support a wider range of layouts when dealing with unpacking.
    # 16-bit types are special, because the store instruction can unpack them.
    elif (
        value.layout == TMEM_NATIVE_LAYOUT
        and has_default_layout
        and ((bitwidth == 16 and packing == 1) or bitwidth == 32)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Do not store into sparse metadata allocations; populate them only via async_copy_sparse_smem_to_tmem
2. Allocate a separate TMEM buffer for values you need to store
3. Tag sparse allocations in your kernel code so store is never called on them

if tmem.layout == tcgen05.sparse_meta_layout():
    raise NotImplementedError('cannot store to sparse metadata TMEM')

def is_sparse_meta(tmem) -> bool:
    return tmem.layout == tcgen05.sparse_meta_layout()

When it happens

Trigger: Allocating TMEM with sparse_meta_layout() for sparse MMA and later calling tmem.store(value) on it.

Common situations: Reusing a TMEM allocation for both sparse metadata and general storage; trying to initialize sparse metadata from registers.

Related errors

• Sparse meta layout loads unsupported.
• Loading multiple row tiles
• Sparse MMA not supported for M=64
• Sparse MMA unsupported for f32
• Loads from TMEM layout {self.layout} to register layout {lay

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