jax-ml/jax · error · ValueError

Base indices of array slices must be aligned to the beginnin

Error message

Base indices of array slices must be aligned to the beginning of a tile. The array uses a tiling of {base_tile_shape}, but your base indices are {base_idx}. Consider using a different array layout.

What it means

When slicing a tiled FragmentedArray, every base index must be a multiple of the corresponding tile dimension, otherwise ValueError is raised telling you the tiling and your base indices. Register-level slicing can only select whole tiles.

Source

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

    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(
        b if sq else slice(b // t, (b + l) // t)
        for b, l, t, sq in zip(
            base_idx, slice_shape, base_tile_shape, is_squeezed, strict=True
        )
    )
    new_regs = self.registers[register_slices]
    return FragmentedArray(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align slice start offsets to multiples of the tile shape (round the start down/up to a tile boundary)
2. Choose a base_tile_shape that divides your chunk size
3. Relayout to a tile shape compatible with your slicing pattern via to_layout

Example fix

# before
sub = fa[20:52, :]  # tile size 8 -> 20 % 8 != 0
# after
sub = fa[24:56, :]  # aligned to tile boundaries of 8

tile = fa.layout.base_tile_shape
assert all(b % t == 0 for b, t in zip(base_idx, tile)), 'start not tile-aligned'

When it happens

Trigger: fa[4:68, :] where dim 0 has tile size 8: 4 % 8 != 0, so the slice start is mid-tile.

Common situations: Assuming numpy-style arbitrary-range slicing on tile-blocked fragments, e.g. chunked epilogues with offsets not aligned to the tile shape.

Related errors

• The slice shape must be a multiple of the tile shape. The ar
• Only arrays with tiled layouts can be sliced
• Only slicing with static indices allowed
• Integer indexing not implemented for tiled dimensions (only
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp

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