jax-ml/jax · error · ValueError

Swizzle {self.swizzle} is not supported. Only 32, 64 and 128

Error message

Swizzle {self.swizzle} is not supported. Only 32, 64 and 128 are accepted.

What it means

SwizzleTransform only accepts swizzle values 32, 64, or 128 (byte widths understood by TPU shared-memory swizzling hardware); anything else fails validation in __post_init__.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1173

    return cls(dtypes.dtype(ref.dtype), ref.ref.shape, byte_offset, alias_group_idx, layout)

  def transform_type(self, x):
    match x:
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case jax_core.ShapedArray():
        return x.update(shape=self.shape, dtype=self.dtype)
      case _:
        raise TypeError(f"Unsupported type: {x}")


@dataclasses.dataclass(frozen=True)
class SwizzleTransform(state_types.Transform):
  swizzle: int

  def __post_init__(self):
    if self.swizzle not in {32, 64, 128}:
      raise ValueError(
          f"Swizzle {self.swizzle} is not supported. Only 32, 64 and 128 are"
          " accepted."
      )

  def transform_type(
      self, x: jax_core.AbstractValue
  ) -> jax_core.AbstractValue:
    match x:
      case jax_core.ShapedArray():
        swizzle_elems = (self.swizzle * 8) // dtypes.itemsize_bits(x.dtype)
        if swizzle_elems != x.shape[-1]:
          raise ValueError(
              f"Swizzle {self.swizzle} requires the trailing dimension to be of"
              f" size {swizzle_elems}, but got shape: {x.shape}"
          )
        return x
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp/choose swizzle from {32, 64, 128} only
2. Derive swizzle from the TMA block layout (typically 128) rather than computing arbitrary values
3. Validate user-supplied swizzle before constructing the transform

Example fix

# before
SwizzleTransform(swizzle=16)

# after
assert swizzle in (32, 64, 128)
SwizzleTransform(swizzle=swizzle)

if swizzle not in (32, 64, 128): raise ValueError('swizzle must be 32, 64 or 128')

def valid_swizzle(s): return s in (32, 64, 128)

When it happens

Trigger: Constructing SwizzleTransform(swizzle=16) or other invalid values, usually by parameterizing swizzle from user config in a pallas kernel.

Common situations: Passing a swizzle derived from dtype size or block shape arithmetic that lands outside {32, 64, 128}; copying example code with a wrong constant.

Understand the failure class

Background: "Invalid value" and "allowed values are" config errors: what your library rejected and how to fix it — this error's family across 41 libraries.

Related errors

• dims and idxs must have the same length
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't transpose the swizzled dimension.
• Reshape shape {shape} is not divisible by swizzle elements {
• cluster= must be at most 3D, got {self}.

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