jax-ml/jax · error · ValueError

Loads from TMEM layout {self.layout} to register layout {lay

Error message

Loads from TMEM layout {self.layout} to register layout {layout} are not supported

What it means

TensorMem.load fell through all supported layout cases (default layout, native tiled, wgmma half-lane, m64 collective) and none matched the (self.layout, layout) pair. This is the catch-all for unsupported TMEM-to-register layout conversions.

Source

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

      )
      assert raw_registers.shape[0] == 4
      registers = np.concatenate([raw_registers[:2], raw_registers[2:]], axis=1)
      registers = registers.T.reshape(regs_shape)
    elif (
        layout == fa_m64_collective_layout(columns)
        and self.layout == tmem_m64_collective_layout(columns, packing)
        and is_at_least_16b
    ):
      if reduce is not None:
        raise ValueError("Fused load-reduce is not supported for this layout")
      reduced_reg = None
      regs_shape = layout.registers_shape(self.shape)
      # We take half the columns, because they are split over halves of TMEM.
      registers = _load_32xcols(
          self.address, columns // 2, self.dtype, packing
      ).reshape(regs_shape)
    else:
      raise ValueError(
          f"Loads from TMEM layout {self.layout} to register layout"
          f" {layout} are not supported"
      )
    result = fa.FragmentedArray(
        _registers=registers, _layout=layout, _is_signed=is_signed
    )
    if reduce is None:
      # The None assignments in the branches let us use the linter to ensure
      # that we didn't forget to handle reduce in any of the cases.
      assert reduced_reg is None
      return result
    reduced_layout = layout.reduce((len(layout.base_tile_shape) - 1,))
    assert reduced_layout.vector_length == 1
    reduced_regs_shape = reduced_layout.registers_shape(self.shape[:-1])
    assert math.prod(reduced_regs_shape) == 1
    reduced_result = fa.FragmentedArray(
        _registers=np.asarray(reduced_reg, dtype=object).reshape(reduced_regs_shape),
        _layout=reduced_layout,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Read the branches above the raise to pick a supported layout pair (LAYOUT + default TMEM layout, as_tiled_layout with packing*bitwidth==32, WGMMA_LAYOUT + half-lane, m64 collective)
2. Ensure packing satisfies 32-bit column constraints (e.g. packing = 32 // bitwidth)
3. Construct TMEM with a standard layout via tmem_default_layout / tmem_half_lane_layout / tmem_m64_collective_layout helpers

def load_pair_supported(tmem, layout, bitwidth):
    dl = tcgen05.tmem_default_layout(tmem.packing)
    return ((layout == tcgen05.LAYOUT and tmem.layout == dl and bitwidth in (16,32))
        or (layout == tmem.layout.as_tiled_layout() and tmem.packing*bitwidth == 32)
        or (layout == fa.WGMMA_LAYOUT and tmem.layout == tcgen05.tmem_half_lane_layout(tmem.shape[1], tmem.packing))
        or (layout == tcgen05.fa_m64_collective_layout(tmem.shape[1]) and tmem.layout == tcgen05.tmem_m64_collective_layout(tmem.shape[1], tmem.packing)))

try:
    arr, red = tmem.load(layout, reduce=reduce)
except ValueError as e:
    if 'not supported' in str(e):
        layout = choose_supported_layout(tmem); arr, _ = tmem.load(layout)
    else: raise

When it happens

Trigger: Calling tmem.load(layout) with a register layout that doesn't match any supported pair for the TMEM's layout — e.g. a custom TiledLayout, a transposed layout, or an unsupported packing/bitwidth combination (packing*bitwidth != 32 on the native path).

Common situations: Custom register layouts for new kernels; 8-bit or 64-bit element types where packing constraints fail; layout mismatch after kernel refactors.

Related errors

• Loading multiple row tiles
• Stored array has layout {value.layout}, but TMEM stores expe
• Storing from register layout {value.layout} to TMEM layout {
• Sparse meta layout loads unsupported.
• Fused load-reduce is not supported for this layout

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