jax-ml/jax · error · ValueError

Storing from register layout {value.layout} to TMEM layout {

Error message

Storing from register layout {value.layout} to TMEM layout {self.layout} is not supported

What it means

Catch-all in TensorMem.store: the (value.layout, self.layout) pair matched none of the supported store paths (default layout, native tiled, wgmma half-lane, m64 collective). The register layout cannot be lowered into this TMEM layout.

Source

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

      _store_32xcols_native(self.address, value.registers.reshape(-1), packing)
    elif (
        value.layout == fa.WGMMA_LAYOUT
        and self.layout == tmem_half_lane_layout(self.shape[1], packing=packing)
        and is_at_least_16b
    ):
      registers = value.registers.T.reshape(2, -1)
      registers = np.concatenate(np.split(registers, 2, axis=1), axis=0)
      _store_32xcols(self.address, registers, packing)
    elif (
        value.layout == fa_m64_collective_layout(self.shape[1])
        and self.layout == tmem_m64_collective_layout(
            self.shape[1], packing=packing
        )
        and is_at_least_16b
    ):
      _store_32xcols(self.address, value.registers.reshape(4, -1), packing)
    else:
      raise ValueError(
          f"Storing from register layout {value.layout} to TMEM layout"
          f" {self.layout} is not supported"
      )

  def _debug_print(self) -> None:
    i32 = ir.IntegerType.get_signless(32)
    num_cols = self.layout.cols_in_shape(self.shape, utils.bitwidth(self.dtype))
    lane = arith.remui(utils.thread_idx(), arith.constant(i32, utils.WARPGROUP_SIZE))
    for c in range(num_cols):
      ptr = _tmem_addr_to_ptr(arith.addi(self.address, arith.constant(i32, c)))
      i32_vec = ir.VectorType.get((1,), i32)
      vec_val = nvvm.tcgen05_ld(i32_vec, nvvm.Tcgen05LdStShape.SHAPE_32X32B, ptr)
      val = vector.extract(vec_val, [], [0])
      dtype_bitwidth = utils.bitwidth(self.dtype)
      full_packing = 32 // dtype_bitwidth
      if self.packing == 1:
        if dtype_bitwidth < 32:
          val = arith.trunci(ir.IntegerType.get_signless(dtype_bitwidth), val)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Relayout the value to match one of the supported pairs (e.g. tmem_default_layout with LAYOUT registers)
2. Use a standard TMEM layout helper for allocation
3. Ensure element bitwidth is 16 or 32 and packing satisfies the 32-bit column constraint

assert utils.bitwidth(tmem.dtype) in (16, 32), 'store needs >=16b elements'
assert value.layout == tmem.layout.as_tiled_layout() or tmem.layout == tcgen05.tmem_default_layout(packing=tmem.packing), 'unsupported store pair'

try:
    tmem.store(value)
except ValueError:
    value = value.relayout(tmem.layout.as_tiled_layout()); tmem.store(value)

When it happens

Trigger: Storing a FragmentedArray whose TiledLayout doesn't pair with the TMEM allocation's layout, or with unsupported packing/bitwidth (e.g. 8-bit elements where is_at_least_16b fails).

Common situations: Custom tiled layouts; sub-16-bit element types; layout drift after refactoring a kernel's TMEM allocation strategy.

Related errors

• Loading multiple row tiles
• Loads from TMEM layout {self.layout} to register layout {lay
• Stored array has layout {value.layout}, but TMEM stores expe
• 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/167bd6a1d9830052. Report an issue: GitHub.