{"record":{"id":"14e8e560576ec5c9","repo":"jax-ml/jax","slug":"non-slice-indices-are-not-supported-in-2-minormost","errorCode":null,"errorMessage":"Non-slice indices are not supported in 2 minormost dims: {idxs}","messagePattern":"Non-slice indices are not supported in 2 minormost dims: (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1249,"sourceCode":"          f\" {self.swizzle_elems(aval.dtype)}\"\n      )\n    return transform, self\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:\n    if not hasattr(aval, \"dtype\"):\n      raise ValueError(\n          f\"Cannot commute unswizzle and indexer with {aval}, which does not\"\n          \" have a dtype\"\n      )\n    dtype = aval.dtype\n    swizzle_elems = self.swizzle_elems(dtype)\n    idxs = indexer.indices\n    if not idxs:\n      return indexer, self\n    if not all(isinstance(idx, (slice, indexing.Slice)) for idx in idxs[-2:]):\n      raise NotImplementedError(\n          f\"Non-slice indices are not supported in 2 minormost dims: {idxs}\"\n      )\n    last_idx = idxs[-1]\n    if isinstance(last_idx, indexing.Slice):\n      if last_idx.start != 0 or last_idx.size != swizzle_elems:\n        raise ValueError(\"Swizzled dims cannot be sliced\")\n    else:\n      assert isinstance(last_idx, slice)\n      if (\n          (last_idx.step is not None and last_idx.step != 1)\n          or (last_idx.start is not None and last_idx.start != 0)\n          or (last_idx.stop is not None and last_idx.stop != swizzle_elems)\n      ):\n        raise ValueError(\"Swizzled dims cannot be sliced\")\n    return indexer, self\n\n  def pretty_print(self, context: jax_core.JaxprPpContext) -> pp.Doc:\n    return pp.text(f\"{{unswizzle({self.swizzle})}}\")","sourceCodeStart":1231,"sourceCodeEnd":1267,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1231-L1267","documentation":"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.","triggerScenarios":"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.","commonSituations":"Gathering rows/columns from a swizzled WGMMA buffer inside a Mosaic kernel; adapting normal-layout kernel code that used integer indexing to swizzled layouts.","solutions":["Convert integer indices in the two minormost dims into full-dim slices and select within the block afterwards","Unswizzle the ref before integer indexing","Restructure so integer indexing happens on leading (batch) dims only"],"exampleFix":"// before\nrow = ref[:, 3]  # int index in minormost dim -> NotImplementedError\n// after\nblk = unswizzle(ref)\nrow = blk[:, 3]","handlingStrategy":"validation","validationCode":"assert all(isinstance(i, (slice, indexing.Slice)) for i in idxs[-2:]), 'two minormost dims of swizzled ref must use slices'","typeGuard":null,"tryCatchPattern":"try:\n    ref[idxs]\nexcept NotImplementedError:\n    unswizzle(ref)[idxs]","preventionTips":["Use slices (full swizzle groups) on minormost dims","Unswizzle before integer indexing"],"tags":["jax","pallas","indexing","swizzle","not-implemented"],"backgroundTag":"unsupported-indexing-on-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}