jax-ml/jax · error · ValueError

TMEM layout {self.layout} is not supported

Error message

TMEM layout {self.layout} is not supported

What it means

tcgen05.load converts TMEM to registers using a fragment-array layout matched to the TMEM layout and element bitwidth. Only specific combos are supported: WGMMA_LAYOUT, the m64 collective pair, or the layout as-is when packing*bitwidth==32. Any other TMEM layout raises 'TMEM layout … is not supported'.

Source

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

      layout: fa.TiledLayout | None = None,
      is_signed: bool | None = None,
      reduce: LoadReduceOp | None = None,  # Reduction operator for the minor dimension.
  ) -> fa.FragmentedArray | tuple[fa.FragmentedArray, fa.FragmentedArray]:
    packing = self.packing
    bitwidth = utils.bitwidth(self.dtype)
    is_at_least_16b = bitwidth in {16, 32}
    columns = self.shape[1]
    if layout is None:
      if is_at_least_16b and self.layout == tmem_default_layout(packing):
        layout = LAYOUT
      elif is_at_least_16b and packing <= columns // 2 and self.layout == tmem_half_lane_layout(columns, packing):
        layout = fa.WGMMA_LAYOUT
      elif is_at_least_16b and columns % 16 == 0 and self.layout == tmem_m64_collective_layout(columns, packing):
        layout = fa_m64_collective_layout(columns)
      elif packing * bitwidth == 32:
        layout = self.layout.as_tiled_layout()
      else:
        raise ValueError(f"TMEM layout {self.layout} is not supported")
    if reduce is not None:
      if isinstance(self.dtype, ir.IntegerType) and bitwidth == 32:
        if reduce not in ("min", "max"):
          raise ValueError(
              "Unsupported reduction for i32. Only min and max are supported,"
              f" got: {reduce}"
          )
        if not is_signed:
          reduce = "abs" + reduce  # type: ignore
      elif isinstance(self.dtype, ir.F32Type):
        if reduce not in ("min", "max", "absmin", "absmax"):
          raise ValueError(
              "Unsupported reduction for f32. Only min, max, absmin, and"
              f" absmax are supported, got: {reduce}"
          )
      else:
        raise ValueError(f"Unsupported dtype for reduction: {self.dtype}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the canonical layouts: tmem_default_layout, tmem_half_lane_layout, or tmem_m64_collective_layout
2. Choose packing so packing * bitwidth == 32 (f32->1, f16/bf16->2, 8-bit->4, 4-bit->8) or one of the other supported branches
3. If you need an unsupported layout, relayout in registers after a supported load

Example fix

# before
ref = tcgen05.TMEMRef(..., layout=custom_layout)  # packing*bitwidth != 32
out = tcgen05.load(ref)
# after
ref = tcgen05.TMEMRef(..., layout=tcgen05.tmem_default_layout(packing=2))  # bf16: 2*16==32
out = tcgen05.load(ref)

bitwidth = ref.dtype.width
packing = ref.layout.vector_size  # or however packing is tracked
assert packing * bitwidth == 32 or ref.layout in (known_supported,)

def load_supported(ref) -> bool:
    bitwidth = getattr(ref.dtype, 'width', 32)
    return ref.layout in SUPPORTED_TMEM_LAYOUTS_FOR(ref) or ref.layout_packing * bitwidth == 32

try:
    out = tcgen05.load(ref)
except ValueError as e:
    if 'not supported' in str(e):
        ref = ref.to_layout(tcgen05.tmem_default_layout(packing=32 // bitwidth))
        out = tcgen05.load(ref)
    else:
        raise

When it happens

Trigger: Calling tcgen05.load on a TMEM ref with a custom/relayouted TMEMLayout that doesn't match the supported patterns, or where packing * element_bitwidth != 32 (e.g. packing=1 with f32, or packing=2 with i8).

Common situations: Building custom TMEMLayouts instead of the provided constructors; using packing values inconsistent with dtype width (e.g. packing=4 with 16-bit types gives 64 != 32).

Related errors

• Cannot slice TMEM with multiple tiles along rows.
• Minor dimension of shape must be divisible by packing, got:
• Packing must be a power of 2, got: {packing}
• Columns must be a multiple of 16, got: {columns}
• TMEM can only be sliced, not indexed

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