jax-ml/jax · error · IndexError

Slice {idx} along axis {axis} is out of bounds for shape {sh

Error message

Slice {idx} along axis {axis} is out of bounds for shape {shape}

What it means

Slice bounds on a memref axis are validated: after normalizing negative start/stop, the slice must satisfy 0 <= start < bound and 0 <= stop <= bound. Otherwise Mosaic raises IndexError with the offending slice, like numpy would.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:976

        raise IndexError(
            f"Index {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(idx if idx >= 0 else bound + idx)
      slice_shape.append(1)
      is_squeezed.append(True)
    elif isinstance(idx, slice):
      if idx.step is not None and idx.step != 1:
        raise NotImplementedError("Strided slices not implemented")
      start = idx.start or 0
      if start < 0:
        start = bound + start
      stop = idx.stop or bound
      if stop < 0:
        stop = bound + stop
      if check_oob and (
          start < 0 or start >= bound or stop < 0 or stop > bound
      ):
        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(start)
      slice_shape.append(stop - start)
      is_squeezed.append(False)
    elif isinstance(idx, DynamicSlice):
      if check_oob and (
          isinstance(idx.base, int) and idx.base + idx.length > bound
      ):
        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(idx.base)
      slice_shape.append(idx.length)
      is_squeezed.append(False)
    elif isinstance(idx, ir.Value):
      if not isinstance(idx.type, ir.IndexType):
        raise ValueError("Expected an index-typed index")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp the slice: stop = min(stop, bound); start = max(0, min(start, bound)) before indexing
2. Pad the source tensor so slice extents divide evenly into tiles
3. Compute slices from the actual shape: ir.MemRefType(ref.type).shape[axis]

Example fix

# before
edge = buf[off : off + BLOCK]  # off+BLOCK > size on last tile
# after
size = ir.MemRefType(buf.type).shape[0]
edge = buf[off : min(off + BLOCK, size)]

bound = ir.MemRefType(ref.type).shape[axis]
start = max(0, s.start or 0)
stop = max(start, min(s.stop if s.stop is not None else bound, bound))
assert 0 <= start < bound and start <= stop <= bound

When it happens

Trigger: ref[10:20] on an axis of size 8; ref[5:-10] where -10 normalizes below 0; ref[start:] with start computed from a loop variable exceeding the bound; with check_oob enabled.

Common situations: Edge tiles in a tiled kernel where the tile slice extends past the tensor extent; reusing slice parameters after changing the buffer shape; forgetting that unlike numpy, Mosaic does not silently clamp stop to bound.

Related errors

• Folding {fold_rank} dimensions starting from {dim} is out of
• Index {idx} along axis {axis} is out of bounds for shape {sh
• Strided slices not implemented
• SubViewOp only supports a single tile transform.
• Only support memref.cast where the input and output types ar

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