jax-ml/jax · error · ValueError

TMEM can only be sliced, not indexed

Error message

TMEM can only be sliced, not indexed

What it means

TMEM slicing produces a view, not a copy, so removing a dimension (indexing with an integer, which squeezes the axis) has no meaning for TMEM addressing. utils.parse_indices marks squeezed dims, and TMEMRef.slice rejects any non-None integer index.

Source

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

    if shape[0] < 32:
      raise ValueError(f"TMEM refs must have at least 32 rows, got: {shape[0]}")
    if layout is None:
      if collective is None:
        raise ValueError(
            "collective argument must be provided when TMEM layout is inferred"
        )
      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]):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use slices with explicit ranges: tmem_ref[0:shape0, a:b]
2. Select columns in registers after tcgen05.load instead of indexing TMEM
3. When lowering subviews, keep both dims (use 1-wide slices, but note row slicing limits)

Example fix

# before
sub = tmem_ref[3]  # indexing
# after
sub = tmem_ref[:, 3:4]  # slicing (columns only; rows must stay full)

import numpy as np
idxs = (slice(0, shape[0]), slice(a, b))
assert all(isinstance(i, slice) for i in idxs), 'TMEM supports slicing only'

def is_slice_only(idxs) -> bool:
    return all(i is None or isinstance(i, slice) for i in idxs)

When it happens

Trigger: tmem_ref[0], tmem_ref[:, 3], or tmem_ref[5, :] — any integer index; also reached from apply_fun, broadcast_into, or lowering rules for subview/insert_strided_slice ops that index the TMEM ref.

Common situations: Writing numpy-style indexing on TMEM refs; MLIR-level subview lowering that fully reduces a dimension; converting register code that indexed accumulators.

Related errors

• TMEM cannot be sliced along rows
• 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/93a20c724037a254. Report an issue: GitHub.