jax-ml/jax · error · NotImplementedError

Only arrays with tiled layouts can be sliced

Error message

Only arrays with tiled layouts can be sliced

What it means

FragmentedArray.__getitem__ (slicing) only supports TiledLayout (and WGSplatFragLayout, which is handled earlier); slicing any other layout, such as a slice of a sliced array with non-tiled layout, raises NotImplementedError('Only arrays with tiled layouts can be sliced').

Source

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

      raise ValueError("Only bitcast between types of the same bitwidth supported")
    reg_type = self.registers.flat[0].type
    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."
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call fa.to_layout(TiledLayout(...)) to relayout before slicing
2. Slice at the JAX/lax level before lowering to Mosaic
3. Avoid slicing fragments with non-standard layouts; restructure indices instead

Example fix

# before
sub = fa[0:128, :]
# after
from jax.experimental.mosaic.gpu.layout import TiledLayout
sub = fa.to_layout(TiledLayout((128, 8)))[0:128, :]

from jax.experimental.mosaic.gpu.layout import TiledLayout, WGSplatFragLayout
assert isinstance(fa.layout, (TiledLayout, WGSplatFragLayout)), type(fa.layout)

def sliceable(fa) -> bool:
    from jax.experimental.mosaic.gpu.layout import TiledLayout, WGSplatFragLayout
    return isinstance(fa.layout, (TiledLayout, WGSplatFragLayout))

When it happens

Trigger: fa[some_slice] where fa.layout is not TiledLayout — e.g. arrays produced by operations that yield an untiled/weird layout, or nested slicing on an already-sliced fragment that changed layout.

Common situations: Trying to slice a fragment that came from to_layout to a non-tiled layout, or a custom/opaque layout after transformations.

Related errors

• {op} has an unsupported layout: {out_layout_attr}
• Expected TiledLayout, got {type(layout)}
• Output layout {out_layout} must match the accumulator layout
• Unsupported layout: {src.layout}
• Only slicing with static indices allowed

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