jax-ml/jax · error · NotImplementedError

Loading multiple row tiles

Error message

Loading multiple row tiles

What it means

TensorMem.load only supports tensors whose register shape has a single row tile (regs_shape[0] == 1). When the requested register layout needs multiple row tiles, the load would require issuing several loads, which is not implemented.

Source

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

          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
      # dimensions that map to columns map only to columns.
      if reduce is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a register layout whose registers_shape(tmem.shape)[0] == 1 (e.g. LAYOUT or wgmma-compatible layouts)
2. Check layout.registers_shape(tmem.shape) before calling load
3. Resize the TMEM allocation or restructure the kernel so one row tile suffices

regs_shape = layout.registers_shape(tmem.shape)
assert regs_shape[0] == 1, f'load supports one row tile, got {regs_shape[0]}'

def load_supported(tmem, layout) -> bool:
    return layout.registers_shape(tmem.shape)[0] == 1

When it happens

Trigger: Calling tmem.load(layout) where layout.registers_shape(tmem.shape)[0] != 1 — e.g. tall/narrow register layouts spanning more than 128 lanes worth of rows for the TMEM allocation.

Common situations: Choosing a FragmentedArray register layout with more row tiles than the TMEM shape supports; mismatched layouts when refactoring kernels from one shape to another.

Related errors

• Sparse meta layout loads unsupported.
• Loads from TMEM layout {self.layout} to register layout {lay
• Stored array has layout {value.layout}, but TMEM stores expe
• Sparse meta layout stores unsupported.
• Storing from register layout {value.layout} to TMEM layout {

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