{"record":{"id":"b465c1b16747d657","repo":"jax-ml/jax","slug":"integer-indexing-not-implemented-for-tiled-dimensi","errorCode":null,"errorMessage":"Integer indexing not implemented for tiled dimensions (only slicing allowed)","messagePattern":"Integer indexing not implemented for tiled dimensions \\(only slicing allowed\\)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":2069,"sourceCode":"\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      )\n    if any(l % t for l, t in zip(slice_shape, base_tile_shape, strict=True)):\n      raise ValueError(\n          \"The slice shape must be a multiple of the tile shape. The array\"\n          f\" uses a tiling of {base_tile_shape}, but your slice shape is\"\n          f\" {slice_shape}. Consider using a different array layout.\"\n      )\n    register_slices = tuple(","sourceCodeStart":2051,"sourceCodeEnd":2087,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L2051-L2087","documentation":"Slicing a tiled FragmentedArray raises NotImplementedError('Integer indexing not implemented for tiled dimensions') when an integer index squeezes one of the tiled dimensions. Only full slices (start:stop) are supported along tiled dims because registers hold whole tiles.","triggerScenarios":"fa[0, 0:64] where dimension 0 is tiled: the integer 0 squeezes a tiled dimension, which cannot drop part of a tile from registers.","commonSituations":"Indexing a single row/column of a tile-blocked tensor like a normal numpy array inside a Mosaic kernel.","solutions":["Keep slices (e.g. 0:1) instead of integer indices on tiled dims and squeeze later","Use a non-tiled layout for dims you need to integer-index","Do fine-grained indexing with loads from memory rather than fragment slicing"],"exampleFix":"# before\nrow = fa[3, :]  # dim 0 is tiled\n# after\nrow = fa[3:4, :]  # slice keeps the dim; handle the size-1 dim downstream","handlingStrategy":"validation","validationCode":"untiled = len(fa.shape) - len(fa.layout.base_tile_shape)\nassert not any(is_squeezed[untiled:]), 'no integer indexing on tiled dims'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use start:stop slices (size-1) instead of ints on tiled dims","Index via loads for single elements"],"tags":["mosaic","gpu","slicing","tiled-layout"],"backgroundTag":"unsupported-indexing-on-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}