jax-ml/jax · error · NotImplementedError

Integer indexing not implemented for tiled dimensions (only

Error message

Integer indexing not implemented for tiled dimensions (only slicing allowed)

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:2069

    return self._pointwise(
        lambda x: arith.bitcast(ty, x), output_is_signed=output_is_signed, restrict_bitwidth=False
    )

  def __getitem__(self, idx) -> FragmentedArray:
    base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape)
    if isinstance(self.layout, WGSplatFragLayout):
      shape = tuple(d for d, s in zip(slice_shape, is_squeezed) if not s)
      return self.splat(self.registers.item(), shape, is_signed=self.is_signed)
    if not isinstance(self.layout, TiledLayout):
      raise NotImplementedError("Only arrays with tiled layouts can be sliced")
    if any(isinstance(idx, ir.Value) for idx in base_idx):
      raise ValueError("Only slicing with static indices allowed")
    base_idx = cast(tuple[int, ...], base_idx)
    base_tile_shape = self.layout.base_tile_shape
    untiled_rank = len(self.shape) - len(base_tile_shape)
    if any(is_squeezed[untiled_rank:]):
      raise NotImplementedError(
          "Integer indexing not implemented for tiled dimensions (only slicing"
          " allowed)"
      )
    if untiled_rank:
      base_tile_shape = (1,) * untiled_rank + base_tile_shape
    if any(b % t for b, t in zip(base_idx, base_tile_shape, strict=True)):
      raise ValueError(
          "Base indices of array slices must be aligned to the beginning of a"
          f" tile. The array uses a tiling of {base_tile_shape}, but your base"
          f" indices are {base_idx}. Consider using a different array layout."
      )
    if any(l % t for l, t in zip(slice_shape, base_tile_shape, strict=True)):
      raise ValueError(
          "The slice shape must be a multiple of the tile shape. The array"
          f" uses a tiling of {base_tile_shape}, but your slice shape is"
          f" {slice_shape}. Consider using a different array layout."
      )
    register_slices = tuple(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep slices (e.g. 0:1) instead of integer indices on tiled dims and squeeze later
2. Use a non-tiled layout for dims you need to integer-index
3. Do fine-grained indexing with loads from memory rather than fragment slicing

Example fix

# before
row = fa[3, :]  # dim 0 is tiled
# after
row = fa[3:4, :]  # slice keeps the dim; handle the size-1 dim downstream

untiled = len(fa.shape) - len(fa.layout.base_tile_shape)
assert not any(is_squeezed[untiled:]), 'no integer indexing on tiled dims'

When it happens

Trigger: 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.

Common situations: Indexing a single row/column of a tile-blocked tensor like a normal numpy array inside a Mosaic kernel.

Related errors

• {op} has an unsupported layout: {out_layout_attr}
• Only arrays with tiled layouts can be sliced
• Only slicing with static indices allowed
• Base indices of array slices must be aligned to the beginnin
• The slice shape must be a multiple of the tile shape. The ar

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