jax-ml/jax · error · ValueError

Swizzle {self.swizzle} requires the trailing dimension to be

Error message

Swizzle {self.swizzle} requires the trailing dimension to be of size {swizzle_elems}, but got shape: {x.shape}

What it means

SwizzleTransform.transform_type checks that the trailing (minormost) dimension size equals the swizzle width in elements (swizzle_bytes*8 / itemsize_bits). If the last dim doesn't match, the swizzled layout would be ill-defined and a ValueError with the expected size is raised.

Source

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

@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))
      case _:
        raise NotImplementedError(f"Unsupported type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return UnswizzleRef(self.swizzle)


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class UnswizzleRef(state_types.Transform):
  swizzle: int = jax.tree.static()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad/reshape the block so the trailing dim equals (swizzle*8)//itemsize_bits
Choose the swizzle that matches your block: for trailing dim N and dtype, pick 32/64/128 so N == swizzle_elems
Avoid applying the swizzle transform to non-matching intermediate shapes; apply it only at the storage layout level

Example fix

# before
# bf16 block shape (..., 16) with SwizzleTransform(128) -> needs 64 elems
ref = ... SwizzleTransform(128)

# after
# pad trailing dim to 64, or use SwizzleTransform(32) for 16 bf16 elems (16*16=256b -> use 32B swizzle)
ref = ... SwizzleTransform(32)  # matches trailing dim 16 for bf16? verify: 32*8/16=16

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp
from jax._src import dtypes
expected = (swizzle * 8) // dtypes.itemsize_bits(dtype)
assert x.shape[-1] == expected, f'trailing dim must be {expected}'

Type guard

def swizzle_matches(swizzle, shape, dtype): return shape[-1] == (swizzle * 8) // dtypes.itemsize_bits(dtype)

Prevention

Match swizzle bytes to block trailing dim and dtype
Pad blocks to swizzle-aligned trailing dims

When it happens

Trigger: Applying SwizzleTransform(128) to a block whose last dimension isn't swizzle_elems long — e.g. fp32 (32-bit) block with trailing dim != 32 for swizzle 128, or padding changing the last dim.

Common situations: Swizzled TMA layouts where block shapes/dtypes don't match the swizzle: bf16 with swizzle 64 needs trailing dim 32; using padded or ragged trailing dimensions.

Related errors

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