{"record":{"id":"31ae27ffcad512ec","repo":"jax-ml/jax","slug":"swizzled-dims-cannot-be-sliced","errorCode":null,"errorMessage":"Swizzled dims cannot be sliced","messagePattern":"Swizzled dims cannot be sliced","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1255,"sourceCode":"  ) -> 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})}}\")\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass CollapseLeadingBatchDimensionsTransform(state_types.Transform):\n  \"\"\"A transform that collapses leading batch dimensions into the minor dimension.","sourceCodeStart":1237,"sourceCodeEnd":1273,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1237-L1273","documentation":"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').","triggerScenarios":"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.","commonSituations":"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.","solutions":["Slice with start=0 and size=swizzle_elems (i.e. the full dim)","Unswizzle the ref before partial slicing","Pad/reshape so the portion you need aligns to whole swizzle groups"],"exampleFix":"// before\npart = ref[:, :, 0:64]  # size 64 != swizzle_elems -> ValueError\n// after\npart = unswizzle(ref)[:, :, 0:64]","handlingStrategy":"validation","validationCode":"se = swizzle_elems(dtype)\nassert last_slice.start == 0 and last_slice.size == se","typeGuard":null,"tryCatchPattern":"try:\n    ref[idx]\nexcept ValueError as e:\n    if 'Swizzled dims cannot be sliced' in str(e):\n        return unswizzle(ref)[idx]\n    raise","preventionTips":["Always slice the swizzled dim as start=0, size=swizzle_elems","Unswizzle for partial slices"],"tags":["jax","pallas","slicing","swizzle"],"backgroundTag":"layout-slicing-restriction","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}