jax-ml/jax · error · ValueError

Only TiledLayouts support swizzling

Error message

Only TiledLayouts support swizzling

What it means

For WGStridedFragLayout, store only supports the 16-byte-swizzle path; any swizzle value other than 16 raises ValueError('Only TiledLayouts support swizzling'). Strided fragment layouts assume 128-bit (16B) aligned vector transactions, and general swizzling modes exist only for TiledLayout.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3654

      atomic: Literal["add", "max", "min", "and", "or", "xor"] | None = None,
  ) -> None:
    index = ir.IndexType.get()
    i64 = ir.IntegerType.get_signless(64)
    if not isinstance(ref.type, ir.MemRefType):
      raise ValueError(ref)
    match self.layout:
      case WGSplatFragLayout():
        if isinstance(ref, utils.MultimemRef):
          raise NotImplementedError("Splat layout does not support multimem")
        if atomic is not None:
          raise NotImplementedError(
              "Atomic stores not supported for splat layout"
          )
        # All values are the same so swizzle does not affect anything here.
        self._store_untiled_splat(ref)
      case WGStridedFragLayout():
        if swizzle != 16:
          raise ValueError("Only TiledLayouts support swizzling")
        assert isinstance(self.layout, WGStridedFragLayout)
        vec_size = self.layout.vec_size
        bitwidth = utils.bitwidth(self.mlir_dtype)
        total_bits = vec_size * bitwidth
        if total_bits % 8 != 0:
          raise NotImplementedError("Vector length should be a multiple of byte size")
        # pyrefly: ignore[bad-argument-type]
        for get, _update, transfer_ref, idx in self.transfer_strided(ref, vec_size):
          if isinstance(transfer_ref, utils.MultimemRef):
            ptr = utils.memref_ptr(utils.memref_slice(transfer_ref.ref, tuple(idx)))
            if atomic is not None:
              self._store_register_atomic(
                  ptr, get(self.registers), atomic, is_smem=False, multimem=True,
              )
            else:
              utils.multimem_store(ptr, get(self.registers))
          elif atomic is not None:
            is_smem = utils.is_smem_ref(transfer_ref)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use swizzle=16 with WGStridedFragLayout, or pass swizzle=None if no swizzle is needed.
2. Switch the array to a TiledLayout (fa.to_layout / build it with TiledLayout) when you need 32/64/128-byte swizzle modes.

Example fix

# before
fa.store(ref, swizzle=32)  # fa has WGStridedFragLayout

# after
fa.store(ref, swizzle=16)
# or: fa.to_layout(TiledLayout(...)).store(ref, swizzle=32)

from jax.experimental.mosaic.gpu.fragmented_array import WGStridedFragLayout
if isinstance(fa.layout, WGStridedFragLayout):
    assert swizzle in (16, None), 'WGStridedFragLayout only supports swizzle=16'

def valid_swizzle(fa, swizzle):
    if isinstance(fa.layout, WGStridedFragLayout):
        return swizzle in (16, None)
    return True

When it happens

Trigger: fa.store(ref, swizzle=32) (or 0/64/128...) on an array with WGStridedFragLayout. swizzle=16 is the only accepted value for this layout; TiledLayout must be used for other swizzle modes.

Common situations: Copying TiledLayout store examples with swizzle=32/64 while the array was built with a warpgroup strided layout; TMA store descriptors with non-16B swizzle combined with strided fragments.

Related errors

• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't transpose the swizzled dimension.
• Reshape shape {shape} is not divisible by swizzle elements {
• No valid out swizzle{what}: minor dimension has {minor_dim_b

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