jax-ml/jax · error · ValueError

Reshape shape {shape} is not divisible by swizzle elements {

Error message

Reshape shape {shape} is not divisible by swizzle elements {self.swizzle_elems(aval.dtype)}

What it means

UnswizzleRef.commute_reshape requires that after a reshape, the minormost dimension of the new shape equals exactly the swizzle element count for the dtype (shape[-1] == swizzle_elems). A reshape to any other trailing extent would break the swizzled 128-bit vector grouping, so it is rejected.

Source

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

  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(
          f"Cannot commute unswizzle and indexer with {aval}, which does not"
          " have a dtype"
      )
    dtype = aval.dtype
    swizzle_elems = self.swizzle_elems(dtype)
    idxs = indexer.indices
    if not idxs:
      return indexer, self

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute swizzle_elems for your dtype (swizzle_bytes*8 // dtype bit-width) and reshape so the last dim equals it
2. Unswizzle the ref first, reshape freely, then re-swizzle
3. Keep the trailing dimension fixed and fold size changes into leading dims instead

Example fix

// before
ref2 = ref.reshape((m, n))  # n != swizzle_elems -> ValueError
// after
se = swizzle_elems(dtype)
ref2 = ref.reshape((m * n // se, se))

se = swizzle_elems(dtype)
assert new_shape[-1] == se, f'last dim must be {se}, got {new_shape[-1]}'

try:
    ref.reshape(new_shape)
except ValueError:
    unswizzle(ref).reshape(new_shape)

When it happens

Trigger: Calling reshape on a ref with a SwizzleTransform where the resulting last dimension is not swizzle_elems(dtype) (e.g. reshaping to trailing dim 64 or 256 when swizzle_elems is 128 for the dtype).

Common situations: Reshaping WGMMA operand/accumulator blocks in Mosaic kernels to feed another matmul; dtype changes that alter swizzle_elems (e.g. switching f32 to bf16) without adjusting the reshape target.

Related errors

• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unfolding dimensions is not supported when commuting an `Un
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128

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