jax-ml/jax · error · ValueError

Unsupported shape: {shape}. TMEM references must have either

Error message

Unsupported shape: {shape}. TMEM references must have either {TMEM_ROWS} or {TMEM_ROWS // 2} rows, but got: {shape[0]}.

What it means

TMEM on Blackwell is organized into 128 lanes (rows); hardware only supports allocations of 128 rows or 64 rows (half-lane). _infer_tmem_layout dispatches on shape[0] being TMEM_ROWS (128) or TMEM_ROWS//2 (64) and rejects any other row count because no TMEM layout can describe it.

Source

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

    )


def _infer_tmem_layout(shape: tuple[int, ...], collective: bool, packing: int) -> TMEMLayout:
  if len(shape) != 2:
    raise ValueError(f"TMEM can only represent 2D shapes, got {shape}")
  if packing > 8 or packing.bit_count() != 1:
    raise ValueError(f"Packing must be <= 8 and a power of 2, got: {packing}")
  if shape[1] % packing:
    raise ValueError(f"Minor dimension of shape must be divisible by packing, got: {shape}")
  if shape[0] == TMEM_ROWS:
    return tmem_default_layout(packing)
  elif shape[0] == TMEM_ROWS // 2:
    if collective:
      return tmem_m64_collective_layout(shape[1], packing)
    else:
      return tmem_half_lane_layout(shape[1], packing)
  else:
    raise ValueError(
        f"Unsupported shape: {shape}. TMEM references must have either"
        f" {TMEM_ROWS} or {TMEM_ROWS // 2} rows, but got {shape[0]}."
    )


def tmem_default_layout(packing: int = 1) -> TMEMLayout:
  """A TMEM layout used for 1CTA MMA with M=128 and 2CTA MMA with M=256."""
  if packing.bit_count() != 1:
    raise ValueError(f"Packing must be a power of 2, got: {packing}")
  return TMEMLayout(
      fa.Tiling(((TMEM_ROWS, packing), (fa.WARP_SIZE, packing))),
      warp_dims=(-4,),
      lane_dims=(-2,),
      vector_dim=-1,
  )


def tmem_half_lane_layout(columns, packing: int = 1) -> TMEMLayout:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/split the computation so the TMEM-resident operand has 128 or 64 rows (e.g. M=256 2CTA -> 128 rows per CTA with collective=True)
2. For small M, keep registers in SMEM/registers rather than TMEM
3. Pass an explicit valid layout with a compatible shape instead of inference

Example fix

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

TMEM_ROWS = 128
assert shape[0] in (TMEM_ROWS, TMEM_ROWS // 2), f'bad TMEM rows: {shape[0]}'

def is_supported_tmem_rows(shape: tuple[int, ...]) -> bool:
    return len(shape) == 2 and shape[0] in (128, 64)

When it happens

Trigger: from_alloc / infer_tmem_layout / mma with an accumulator shape like (32, N), (96, N), (256, N) — any row count other than 128 or 64. Note from_alloc separately requires >=32 rows, so values like 32 or 96 reach this error.

Common situations: Porting WGMMA (Hopper) kernels that used M=32/48/96 to tcgen05; computing TMEM shape from arbitrary MMA M dimensions; multi-CTA code where the per-CTA slice of M is not 64 or 128.

Related errors

• Minor dimension of shape must be divisible by packing, got:
• tmem_addr_ref must contain a single element, got: {addr_ref_
• TMEM refs must have at least 32 rows, got: {shape[0]}
• Not implemented: bitcast 1D
• Not implemented: the 2nd minor dim can not be perfectly pack

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