jax-ml/jax · error · TypeError

Stored array has layout {value.layout}, but TMEM stores expe

Error message

Stored array has layout {value.layout}, but TMEM stores expect a TiledLayout

What it means

TensorMem.store requires the incoming FragmentedArray's layout to be a TiledLayout (register layouts like LAYOUT or WGMMA_LAYOUT are not accepted). Store lowering only knows how to move tiled-layout registers into TMEM.

Source

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

        _is_signed=is_signed,
    )
    return result, reduced_result

  def store(self, value: fa.FragmentedArray):
    if not isinstance(value, fa.FragmentedArray):
      raise TypeError(f"TMEM stores expect a FragmentedArray, got: {value}")
    if value.shape != self.shape:
      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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Relayout the array first: value = value.relayout(some_tiled_layout) so value.layout is a fa.TiledLayout
2. Store via SMEM (round-trip through shared memory) if a suitable tiled layout is unavailable
3. Check isinstance(value.layout, fa.TiledLayout) before storing

Example fix

// before
tmem.store(mma_result)  # layout is WGMMA_LAYOUT
// after
tiled = mma_result.relayout(tmem.layout.as_tiled_layout())
tmem.store(tiled)

if not isinstance(value.layout, fa.TiledLayout):
    value = value.relayout(tmem.layout.as_tiled_layout())

def has_tiled_layout(v) -> bool:
    return isinstance(v.layout, fa.TiledLayout)

try:
    tmem.store(value)
except TypeError:
    tmem.store(value.relayout(tmem.layout.as_tiled_layout()))

When it happens

Trigger: tmem.store(value) where value.layout is a strided register layout such as tcgen05.LAYOUT or fa.WGMMA_LAYOUT instead of a TiledLayout instance.

Common situations: Storing a freshly computed mma result (which has a register layout) back to TMEM; re-layouting via relayout to a tiled layout before store is required.

Related errors

• Loading multiple row tiles
• Loads from TMEM layout {self.layout} to register layout {lay
• TMEM stores expect a FragmentedArray, got: {value}
• Storing from register layout {value.layout} to TMEM layout {
• Sparse meta layout loads unsupported.

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