jax-ml/jax · error · ValueError

Can't transpose the swizzled dimension.

Error message

Can't transpose the swizzled dimension.

What it means

UnswizzleRef.commute_transpose raises when asked to commute a transpose past an unswizzle transform whose permutation moves the last (swizzled/minormost) dimension. The unswizzle layout is only defined while the swizzled dim stays in the minormost position, so transposing it is not expressible.

Source

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

  def transform_type(self, x: jax_core.AbstractValue) -> jax_core.AbstractValue:
    # Swizzling preserves the type
    return x

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

  def swizzle_elems(self, dtype: jax.typing.DTypeLike | ir.Type) -> int:
    if not isinstance(dtype, ir.Type):
      dtype = mgpu_utils.dtype_to_ir_type(dtype)
    return (self.swizzle * 8) // mgpu.bitwidth(dtype)

  def commute_transpose(
      self, _: jax_core.AbstractValue, transpose: state_types.TransposeTransform
  ) -> tuple[state_types.TransposeTransform, UnswizzleRef]:
    perm = transpose.permutation
    if perm[-1] != len(perm) - 1:
      raise ValueError("Can't transpose the swizzled dimension.")
    return transpose, self

  def commute_reshape(
      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform
  ) -> tuple[state_types.ReshapeTransform, UnswizzleRef]:
    shape = aval.shape
    if shape[-1] != self.swizzle_elems(aval.dtype):
      raise ValueError(
          f"Reshape shape {shape} is not divisible by swizzle elements"
          f" {self.swizzle_elems(aval.dtype)}"
      )
    return transform, self

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:
    if not hasattr(aval, "dtype"):
      raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restrict transposes to the leading (non-swizzled) dimensions only
2. Unswizzle (materialize to normal layout) before transposing, then re-swizzle
3. Redesign the kernel so data is stored pre-transposed instead of transposing in-kernel

Example fix

// before
x_t = lax.transpose(x_swizzled, (1, 0))  # perm[-1] != 1 -> ValueError
// after
x_n = unswizzle(x_swizzled)
x_t = lax.transpose(x_n, (1, 0))

def check_transpose_perm(perm):
    assert perm[-1] == len(perm) - 1, 'Cannot transpose the swizzled (minormost) dim'

try:
    y = lax.transpose(x, perm)
except ValueError as e:
    if 'swizzled dimension' in str(e):
        x = unswizzle(x); y = lax.transpose(x, perm)
    else:
        raise

When it happens

Trigger: Calling .T or lax.transpose / doing block transposes on a ref that carries a SwizzleTransform, with a permutation where perm[-1] != len(perm) - 1 (last dim is not kept in place).

Common situations: Writing a Pallas/Mosaic kernel that transposes a swizzled TMA/WGMMA buffer, or composing ops (reshape+transpose) on tpu swizzled layouts.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Can't transpose a TMEM reference.
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Swizzle {self.swizzle} requires the trailing dimension to be
• Reshape shape {shape} is not divisible by swizzle elements {

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