jax-ml/jax · error · NotImplementedError

Sparse meta layout loads unsupported.

Error message

Sparse meta layout loads unsupported.

What it means

TensorMem.load refuses to load from TMEM carrying the sparse metadata layout. Sparse metadata written by async_copy_sparse_smem_to_tmem has a physical TMEM layout that TiledLayout cannot faithfully describe (it would need multiple vector dims), so loading it back through registers would give inconsistent results.

Source

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

      elif isinstance(self.dtype, ir.F32Type):
        if reduce not in ("min", "max", "absmin", "absmax"):
          raise ValueError(
              "Unsupported reduction for f32. Only min, max, absmin, and"
              f" absmax are supported, got: {reduce}"
          )
      else:
        raise ValueError(f"Unsupported dtype for reduction: {self.dtype}")

    has_default_layout = self.layout == tmem_default_layout(packing)
    regs_shape = layout.registers_shape(self.shape)
    # 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 loads unsupported.")
    if regs_shape[0] != 1:  # We'll need to issue multiple loads below.
      raise NotImplementedError("Loading multiple row tiles")
    if (
        layout == LAYOUT
        and self.layout == tmem_default_layout(packing)
        and is_at_least_16b
    ):
      if reduce is not None:
        raise ValueError(
            "Fused load-reduce is not supported for this layout"
        )
      reduced_reg = None
      registers = _load_32xcols(
          self.address, columns, self.dtype, packing
      ).T.reshape(regs_shape)
    elif layout == self.layout.as_tiled_layout() and packing * bitwidth == 32:
      # TODO(apaszke): We raise NotImplemented here because technically for some
      # layouts this does make sense. I think only for those where all

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Don't load sparse metadata; consume it only via tcgen05.mma with sparse descriptors
2. If you need to inspect values, copy them through a separate non-sparse TMEM allocation
3. Track which allocations hold metadata (e.g. a wrapper class) so load is never called on them

if tmem.layout == tcgen05.sparse_meta_layout():
    raise NotImplementedError('cannot load sparse metadata; consume via tcgen05.mma')

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

When it happens

Trigger: Allocating TMEM with sparse_meta_layout() (used with async_copy_sparse_smem_to_tmem for sparse MMA) and then calling .load(...) on it.

Common situations: Debugging sparse-mma kernels by trying to read sparse metadata back to registers; attempting to inspect/print sparsity metadata via load.

Related errors

• Sparse meta layout stores 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/1602ec8d21704352. Report an issue: GitHub.