What causes "TMEM refs must have at least 32 rows, got: {shape[0]}" in jax-ml/jax?

Use at least 32 rows — practically 64 or 128 to satisfy layout inference

jax-ml/jax · error · ValueError

TMEM refs must have at least 32 rows, got: {shape[0]}

Error message

TMEM refs must have at least 32 rows, got: {shape[0]}

What it means

TMEM is allocated in units of 32 lanes minimum on Blackwell; from_alloc rejects shapes with fewer than 32 rows even though the layout inference only supports 64/128 rows. This is an early guard against degenerate allocations that the hardware cannot make.

Source

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

      shape: tuple[int, int],
      dtype,
      collective: bool | None = None,
      layout: TMEMLayout | None = None,
  ) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(tmem_addr_ref.type, ir.MemRefType):
      raise ValueError(f"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}")
    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)
    if not utils.is_smem_ref(addr_ref_ty):
      raise ValueError(f"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}")
    if addr_ref_ty.element_type != i32:
      raise ValueError(f"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}")
    if math.prod(addr_ref_ty.shape) != 1:
      raise ValueError(f"tmem_addr_ref must contain a single element, got: {addr_ref_ty}")
    i0 = arith.ConstantOp.create_index(0)
    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)
    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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use at least 32 rows — practically 64 or 128 to satisfy layout inference
Keep small matrices in registers/SMEM instead of TMEM
Scale test shapes up to the hardware minimum

Example fix

# before
ref = tcgen05.TMEMRef.from_alloc(alloc, (16, 64), collective=True)
# after
ref = tcgen05.TMEMRef.from_alloc(alloc, (64, 64), collective=True)

Defensive patterns

Strategy: validation

Validate before calling

assert shape[0] >= 32, shape

Type guard

def tmem_rows_valid(shape) -> bool:
    return len(shape) == 2 and shape[0] >= 32

Prevention

Don't prototype with tiny TMEM shapes; hardware minimum is 32 lanes and layouts need 64/128
Keep small operands in registers or SMEM

When it happens

Trigger: TMEMRef.from_alloc(alloc, (16, N), ...) or any shape[0] < 32.

Common situations: Trying small test shapes like (8, 64) or (16, 16) when prototyping tcgen05 kernels; slicing a TMEM ref down below 32 rows (slice itself forbids row slicing, so this mostly occurs at allocation time).

Related errors

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