jax-ml/jax · error · ValueError

Swizzled dims cannot be sliced

Error message

Swizzled dims cannot be sliced

What it means

The last (swizzled) dimension of a swizzled ref can only be indexed as the full swizzle group: an indexing.Slice must have start==0 and size==swizzle_elems. Any narrower slice would split a swizzled vector, so it raises ValueError('Swizzled dims cannot be sliced').

Source

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

  ) -> 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
    if not all(isinstance(idx, (slice, indexing.Slice)) for idx in idxs[-2:]):
      raise NotImplementedError(
          f"Non-slice indices are not supported in 2 minormost dims: {idxs}"
      )
    last_idx = idxs[-1]
    if isinstance(last_idx, indexing.Slice):
      if last_idx.start != 0 or last_idx.size != swizzle_elems:
        raise ValueError("Swizzled dims cannot be sliced")
    else:
      assert isinstance(last_idx, slice)
      if (
          (last_idx.step is not None and last_idx.step != 1)
          or (last_idx.start is not None and last_idx.start != 0)
          or (last_idx.stop is not None and last_idx.stop != swizzle_elems)
      ):
        raise ValueError("Swizzled dims cannot be sliced")
    return indexer, self

  def pretty_print(self, context: jax_core.JaxprPpContext) -> pp.Doc:
    return pp.text(f"{{unswizzle({self.swizzle})}}")


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class CollapseLeadingBatchDimensionsTransform(state_types.Transform):
  """A transform that collapses leading batch dimensions into the minor dimension.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Slice with start=0 and size=swizzle_elems (i.e. the full dim)
2. Unswizzle the ref before partial slicing
3. Pad/reshape so the portion you need aligns to whole swizzle groups

Example fix

// before
part = ref[:, :, 0:64]  # size 64 != swizzle_elems -> ValueError
// after
part = unswizzle(ref)[:, :, 0:64]

se = swizzle_elems(dtype)
assert last_slice.start == 0 and last_slice.size == se

try:
    ref[idx]
except ValueError as e:
    if 'Swizzled dims cannot be sliced' in str(e):
        return unswizzle(ref)[idx]
    raise

When it happens

Trigger: Slicing the minormost dim of a swizzled ref with an indexing.Slice whose start != 0 or size != swizzle_elems(dtype) — e.g. ref[:, :, 0:64] when swizzle_elems is 128.

Common situations: Trying to take half a swizzled vector per block in TPU WGMMA pipelines; slicing layouts written for a different dtype with a different swizzle_elems.

Related errors

• Get only supports slices with stride 1, got {strides}
• Strided slices unsupported. Got stride: {ds.stride}
• Expected slice start ({start}) and slice size ({size}) to be
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Swizzle {self.swizzle} requires the trailing dimension to be

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