jax-ml/jax · error · ValueError

The slice shape must be a multiple of the tile shape. The ar

Error message

The slice shape must be a multiple of the tile shape. The array uses a tiling of {base_tile_shape}, but your slice shape is {slice_shape}. Consider using a different array layout.

What it means

Slicing a tiled FragmentedArray requires the slice length in each dimension to be a multiple of that dimension's tile size; otherwise ValueError reports the tiling vs your slice shape. Partial tiles cannot be represented in the register layout.

Source

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

      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(
        _registers=new_regs, _layout=self.layout, _is_signed=self.is_signed
    )

  def __setitem__(self, idx: object, value: FragmentedArray) -> None:
    if not isinstance(value, FragmentedArray):
      raise ValueError(f"Expected a FragmentedArray, got: {value}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Round the slice extent to a multiple of the tile shape (pad if needed and mask overflow)
2. Pick a tile shape that divides your chunk/problem sizes
3. Slice at whole-tile granularity: shape divisible by base_tile_shape in every dim

Example fix

# before
sub = fa[0:100, :]  # tile 8 -> 100 % 8 != 0
# after
sub = fa[0:104, :]  # 104 = 13*8, mask/pad the extra 4 rows downstream

tile = fa.layout.base_tile_shape
assert all(l % t == 0 for l, t in zip(slice_shape, tile)), 'extent not tile-aligned'

When it happens

Trigger: fa[0:100, :] where dim 0 tile size is 8: 100 % 8 != 0, so the slice shape is not tile-aligned.

Common situations: Splitting work into chunk sizes (e.g. 100) that don't divide evenly by the tile shape chosen for the layout (e.g. 8).

Related errors

• Base indices of array slices must be aligned to the beginnin
• 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/ec2488b4a20dd9d5. Report an issue: GitHub.