jax-ml/jax · error · NotImplementedError

TMEM cannot be sliced along rows

Error message

TMEM cannot be sliced along rows

What it means

TMEM rows map to lane/DPMC structure that is not contiguous for arbitrary row ranges, so a slice must start at row 0 and keep all rows. slice() raises NotImplementedError when base_idx[0] != 0 or slice_shape[0] != full row count.

Source

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

        )
      layout = _infer_tmem_layout(shape, collective, packing=1)
    # TODO: Do we have to do this??
    # warp_idx = utils.warp_idx(sync=False)
    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))
    return cls(tmem_addr, shape, dtype, layout)

  def slice(self, *idxs) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    base_idx, slice_shape, is_squeezed = utils.parse_indices(idxs, self.shape)
    slice_shape = cast(tuple[int, int], tuple(slice_shape))
    if any(is_squeezed):
      raise ValueError("TMEM can only be sliced, not indexed")
    if base_idx == [0] * len(base_idx) and slice_shape == self.shape:
      return self  # Trivial slice
    # If we slice along rows, or attempt to extract several rows, then we may
    # end up with a non-contiguous slice of memory.
    if base_idx[0] != 0 or slice_shape[0] != self.shape[0]:
      raise NotImplementedError("TMEM cannot be sliced along rows")
    # If we attempt to extract non-contiguous tiles, then we will end up with a
    # non-contiguous slice of memory.
    # We check that we have a single tile along rows. Hence slicing along
    # columns produces a contiguous slice of memory.
    if self.shape[0] != self.layout.base_tile_shape[0]:
      raise NotImplementedError(
          "Cannot slice TMEM with multiple tiles along rows."
      )
    col_idx = base_idx[1]
    if not isinstance(col_idx, ir.Value):
      col_idx = arith.constant(i32, col_idx)
    if not utils.is_known_divisible(col_idx, self.layout.base_tile_shape[1]):
      raise NotImplementedError(
          "Slicing along columns is not supported when the column index is not"
          " known to be a multiple of the base tile shape"
      )
    if col_idx.type == ir.IndexType.get():
      col_idx = arith.index_cast(i32, col_idx)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Slice only along columns: tmem_ref[:, a:b]
2. Allocate separate 64-row TMEM refs instead of slicing a 128-row one by rows
3. Move row partitioning into the MMA shape or register-level operations

Example fix

# before
half = tmem_ref[0:64, :]  # row slicing on 128-row ref
# after
half = tcgen05.TMEMRef.from_alloc(alloc64, (64, ncols), collective=True)  # separate 64-row alloc

assert r0 == 0 and r1 == ref.shape[0], 'TMEM row slicing unsupported'

try:
    sub = ref[idx]
except NotImplementedError as e:
    if 'along rows' in str(e):
        raise  # restructure: separate allocs, not row slices
    raise

When it happens

Trigger: tmem_ref[10:, :], tmem_ref[0:64, 8:16] on a 128-row ref, or subview lowerings that shrink the row extent.

Common situations: Splitting an accumulator across warps/CTAs by row; adapting SMEM subview code to TMEM; trying to feed half of an M=128 accumulator into two MMAs.

Related errors

• TMEM can only be sliced, not indexed
• Cannot slice TMEM with multiple tiles along rows.
• TMEM layout {self.layout} is not supported
• Unsupported reduction for i32. Only min and max are supporte
• Minor dimension of shape must be divisible by packing, got:

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