{"record":{"id":"162caa68170f1132","repo":"jax-ml/jax","slug":"only-arrays-with-tiled-layouts-can-be-sliced","errorCode":null,"errorMessage":"Only arrays with tiled layouts can be sliced","messagePattern":"Only arrays with tiled layouts can be sliced","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":2062,"sourceCode":"      raise ValueError(\"Only bitcast between types of the same bitwidth supported\")\n    reg_type = self.registers.flat[0].type\n    if isinstance(reg_type, ir.VectorType):\n      reg_shape = ir.VectorType(reg_type).shape\n      ty = ir.VectorType.get(reg_shape, elt)\n    else:\n      ty = elt\n\n    return self._pointwise(\n        lambda x: arith.bitcast(ty, x), output_is_signed=output_is_signed, restrict_bitwidth=False\n    )\n\n  def __getitem__(self, idx) -> FragmentedArray:\n    base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape)\n    if isinstance(self.layout, WGSplatFragLayout):\n      shape = tuple(d for d, s in zip(slice_shape, is_squeezed) if not s)\n      return self.splat(self.registers.item(), shape, is_signed=self.is_signed)\n    if not isinstance(self.layout, TiledLayout):\n      raise NotImplementedError(\"Only arrays with tiled layouts can be sliced\")\n    if any(isinstance(idx, ir.Value) for idx in base_idx):\n      raise ValueError(\"Only slicing with static indices allowed\")\n    base_idx = cast(tuple[int, ...], base_idx)\n    base_tile_shape = self.layout.base_tile_shape\n    untiled_rank = len(self.shape) - len(base_tile_shape)\n    if any(is_squeezed[untiled_rank:]):\n      raise NotImplementedError(\n          \"Integer indexing not implemented for tiled dimensions (only slicing\"\n          \" allowed)\"\n      )\n    if untiled_rank:\n      base_tile_shape = (1,) * untiled_rank + base_tile_shape\n    if any(b % t for b, t in zip(base_idx, base_tile_shape, strict=True)):\n      raise ValueError(\n          \"Base indices of array slices must be aligned to the beginning of a\"\n          f\" tile. The array uses a tiling of {base_tile_shape}, but your base\"\n          f\" indices are {base_idx}. Consider using a different array layout.\"\n      )","sourceCodeStart":2044,"sourceCodeEnd":2080,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L2044-L2080","documentation":"FragmentedArray.__getitem__ (slicing) only supports TiledLayout (and WGSplatFragLayout, which is handled earlier); slicing any other layout, such as a slice of a sliced array with non-tiled layout, raises NotImplementedError('Only arrays with tiled layouts can be sliced').","triggerScenarios":"fa[some_slice] where fa.layout is not TiledLayout — e.g. arrays produced by operations that yield an untiled/weird layout, or nested slicing on an already-sliced fragment that changed layout.","commonSituations":"Trying to slice a fragment that came from to_layout to a non-tiled layout, or a custom/opaque layout after transformations.","solutions":["Call fa.to_layout(TiledLayout(...)) to relayout before slicing","Slice at the JAX/lax level before lowering to Mosaic","Avoid slicing fragments with non-standard layouts; restructure indices instead"],"exampleFix":"# before\nsub = fa[0:128, :]\n# after\nfrom jax.experimental.mosaic.gpu.layout import TiledLayout\nsub = fa.to_layout(TiledLayout((128, 8)))[0:128, :]","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu.layout import TiledLayout, WGSplatFragLayout\nassert isinstance(fa.layout, (TiledLayout, WGSplatFragLayout)), type(fa.layout)","typeGuard":"def sliceable(fa) -> bool:\n    from jax.experimental.mosaic.gpu.layout import TiledLayout, WGSplatFragLayout\n    return isinstance(fa.layout, (TiledLayout, WGSplatFragLayout))","tryCatchPattern":null,"preventionTips":["Relayout to TiledLayout before slicing","Do slicing in JAX before lowering"],"tags":["mosaic","gpu","slicing","layout"],"backgroundTag":"unsupported-layout-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}