jax-ml/jax · error · ValueError

Only slicing with static indices allowed

Error message

Only slicing with static indices allowed

What it means

When slicing a FragmentedArray with a tiled layout, any index that is a dynamic ir.Value (produced at runtime) triggers ValueError('Only slicing with static indices allowed'). Mosaic slicing needs compile-time-constant base indices to compute register slices.

Source

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

    if isinstance(reg_type, ir.VectorType):
      reg_shape = ir.VectorType(reg_type).shape
      ty = ir.VectorType.get(reg_shape, elt)
    else:
      ty = elt

    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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use static Python-int indices for slicing
2. Compute the address/dynamic access via load/store with computed offsets instead of slicing
3. Hoist the slice out of dynamic control flow so indices are trace-time constants

Example fix

# before
sub = fa[fi, :]  # fi is an ir.Value from loop induction
# after
sub = fa[0, :]  # static index; or gather via memory ops with dynamic offset
ptr = base_ptr + fi * stride  # dynamic addressing at memory level

assert all(not isinstance(i, ir.Value) for i in base_idx), 'static indices only'

def all_static(idx) -> bool:
    import ir
    flat = idx if isinstance(idx, tuple) else (idx,)
    return all(not isinstance(i, ir.Value) for i in flat)

When it happens

Trigger: fa[idx_value, :] where idx_value is an ir.Value (dynamic index computed inside the kernel), rather than a Python int or constant.

Common situations: Using loop-carried runtime offsets (e.g. from a dynamic loop index) as slice positions inside a Mosaic kernel.

Related errors

• Only arrays with tiled layouts can be sliced
• Integer indexing not implemented for tiled dimensions (only
• Base indices of array slices must be aligned to the beginnin
• The slice shape must be a multiple of the tile shape. The ar
• 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/342be4ca534bbefc. Report an issue: GitHub.