jax-ml/jax · error · ValueError

No valid out swizzle{what}: minor dimension has {minor_dim_b

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.

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)

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.