jax-ml/jax · error · ValueError

No valid out swizzle

Error message

No valid out swizzle{what}: minor dimension has {minor_dim_bits} bits, which is not a multiple of 128 (16 bytes)

What it means

find_swizzle picks the largest valid shared-memory swizzle (128/64/32/16 bytes) dividing the minor dimension's bit width. If minor_dim_bits is not a multiple of 128 bits (16 bytes), no swizzle is valid and the error names the offending dimension.

Solutions

Adjust the tile's minor dimension (dtype bits x element count) to a multiple of 128 bits, e.g. ensure N * itemsize_bits % 128 == 0
Use a dtype with a standard bit width (8/16/32 bits) and a matching vector length (e.g. 16 elems of 8-bit = 128 bits)
Pass the `what` description string for a clearer error identifying which operand failed

Example fix

# before: 6 elements of 16-bit = 96 bits
find_swizzle(6 * 16)

# after: 8 elements of 16-bit = 128 bits
find_swizzle(8 * 16)

Defensive patterns

Strategy: validation

Validate before calling

assert minor_dim_bits % 128 == 0, f'{minor_dim_bits=} not a multiple of 128 bits (16 bytes)'

Type guard

def swizzle_possible(minor_dim_bits: int) -> bool:
    return minor_dim_bits % 128 == 0 and minor_dim_bits > 0

Try / catch

try:
    sw = find_swizzle(minor_dim_bits, what='operand A')
except ValueError:
    # round the vector length up to a 128-bit multiple
    minor_dim_bits = -(-minor_dim_bits // 128) * 128
    sw = find_swizzle(minor_dim_bits, what='operand A')

Prevention

Choose tile minor dims as multiples of 128 bits (e.g. 16 x 8-bit or 8 x 16-bit elements)
Compute dtype_bits * vector_length and assert divisibility when picking block shapes
Use standard-width dtypes for MMA operands

When it happens

Trigger: Calling find_swizzle(minor_dim_bits) (directly or via the fused matmul helpers matmul0..matmul5) with a minor dimension whose bit width is not a multiple of 128 — e.g. a 96-bit row, or dtype*elements combos like 8 floats of 12-bit packed types.

Common situations: Choosing block shapes/dtypes in a Pallas GPU matmul so that K or N * itemsize_bits isn't a multiple of 128; using 3- or 6-bit packed types or odd vector lengths with the tcgen05/WGMMA path.

Related errors

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

Appendix: source

Thrown at jax/_src/pallas/mosaic_gpu/helpers.py:249

  applied as ``plgpu`` transform:

    transforms = (
        plgpu.TilingTransform((8, 8 * swizzle // elem_bits)),
        plgpu.SwizzleTransform(swizzle))
    )

  Args:
    minor_dim_bits: The number of bits in the minor (last) dimension of the
      memory region. Usually computed as ``dim_size * jnp.finfo(dtype).bits``.
    what: A string describing the operand for which the swizzle is being
      computed. Improves the error message if specified.
  """
  for swizzle_bytes in (128, 64, 32, 16):
    if minor_dim_bits % (swizzle_bytes * 8) == 0:
      return swizzle_bytes
  if what:
    what = " for " + what
  raise ValueError(
      f"No valid out swizzle{what}: minor dimension has"
      f" {minor_dim_bits} bits, which is not a multiple of 128 (16 bytes)"
  )


def planar_snake(
    lin_idx: jax.Array,
    shape: tuple[int | jax.Array, int | jax.Array],
    minor_dim: int,
    tile_width: int,
):
  """Converts a linear index into an index into shape, trying to optimize locality.

  The "space filling curve" this function computes splits the minor dimension
  into tiles of length ``tile_width``. Every other tile has its major dimension
  inverted, so that the iteration order "snakes around" when going from one tile
  to another.

View on GitHub (pinned to 1e1c6a8fc0)