jax-ml/jax · error · NotImplementedError

Cannot slice TMEM with multiple tiles along rows.

Error message

Cannot slice TMEM with multiple tiles along rows.

What it means

After the row checks, slice() requires the ref's row count to equal the layout's base tile row count — with multiple row tiles, a column slice would produce non-contiguous memory across row tiles. Hence 'Cannot slice TMEM with multiple tiles along rows.' when self.shape[0] != layout.base_tile_shape[0].

Source

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

  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)

    # The code below converts from a logical column index to a physical column
    # index.
    physical_cols_in_tile = self.layout.cols_in_shape(
        cast(tuple[int, int], self.layout.base_tile_shape),
        utils.bitwidth(self.dtype),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid column slicing/relayout on multi-row-tile refs; operate on the full ref
2. Use a layout whose base tile rows equal the ref's rows (e.g. tmem_default_layout for 128 rows)
3. Load to registers first, then slice in the fragment-array domain

Example fix

# before
sub = big_tmem_ref[:, 32:64]  # ref has multiple row tiles
# after
regs = tcgen05.load(big_tmem_ref)  # then slice registers, or use a single-row-tile layout

assert ref.shape[0] == ref.layout.base_tile_shape[0], 'multi-row-tile layout: cannot column-slice'

def single_row_tile(ref) -> bool:
    return ref.shape[0] == ref.layout.base_tile_shape[0]

When it happens

Trigger: Calling slice (directly or via to_layout/apply_fun/subview lowering) on a TMEM ref whose layout tiles rows multiple times, e.g. a 128-row ref with a layout whose base tile covers 64 rows.

Common situations: Relayouting or column-slicing accumulators whose layout was constructed with smaller row tiles; mixing layouts from tmem_m64_collective_layout with full 128-row shapes.

Related errors

• TMEM can only be sliced, not indexed
• TMEM cannot be sliced along rows
• TMEM layout {self.layout} is not supported
• Minor dimension of shape must be divisible by packing, got:
• Packing must be a power of 2, got: {packing}

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