jax-ml/jax · error · NotImplementedError

Non-slice indices are not supported in 2 minormost dims: {id

Error message

Non-slice indices are not supported in 2 minormost dims: {idxs}

What it means

When commuting an unswizzle transform past an NDIndexer, only slice-based (Slice or Python slice) indices are supported in the two minormost dimensions, because integer indexing would break the swizzled vector grouping. Any non-slice index in the last two dims raises NotImplementedError.

Source

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

          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
    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})}}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert integer indices in the two minormost dims into full-dim slices and select within the block afterwards
2. Unswizzle the ref before integer indexing
3. Restructure so integer indexing happens on leading (batch) dims only

Example fix

// before
row = ref[:, 3]  # int index in minormost dim -> NotImplementedError
// after
blk = unswizzle(ref)
row = blk[:, 3]

assert all(isinstance(i, (slice, indexing.Slice)) for i in idxs[-2:]), 'two minormost dims of swizzled ref must use slices'

try:
    ref[idxs]
except NotImplementedError:
    unswizzle(ref)[idxs]

When it happens

Trigger: Indexing a swizzled ref with an integer (or other non-slice) index in either of the last two dimensions, e.g. ref[:, i] or ref[i, j] where i/j are ints.

Common situations: Gathering rows/columns from a swizzled WGMMA buffer inside a Mosaic kernel; adapting normal-layout kernel code that used integer indexing to swizzled layouts.

Related errors

• Cannot commute unswizzle and indexer with {aval}, which does
• Slicing batch dimensions is not supported.
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton

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