jax-ml/jax · error · IndexError

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

Error message

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

What it means

When indexing a Mosaic memref via __getitem__/__setitem__ or memref_slice, integer indices are bounds-checked against the axis extent. An index >= bound (or negative beyond -bound) raises this IndexError, mirroring numpy semantics for GPU memref access in kernels.

Source

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

  )


def parse_indices(
    index, shape: Sequence[int], *, check_oob: bool = True
) -> tuple[list[ir.Value | int], list[int], list[bool]]:
  if not isinstance(index, tuple):
    index = (index,)
  if trailing_dims := len(shape) - len(index):
    index += (slice(None),) * trailing_dims
  base_indices: list[ir.Value | int] = []
  slice_shape = []
  is_squeezed = []
  for axis, (idx, bound) in enumerate(zip(index, shape)):
    if isinstance(idx, (ir.Operation, ir.OpView)):
      idx = idx.result
    if isinstance(idx, int):
      if check_oob and (idx >= bound or (idx < 0 and -idx > bound)):
        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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp or guard indices with min/max against the axis extent before indexing
2. Verify the memref shape with ir.MemRefType(ref.type).shape in debug builds and fix hardcoded indices
3. Pass check_oob=False only if you can prove (e.g. predicated execution) the index is never dereferenced

Example fix

# before
x = tile[row, col]  # col can be >= tile width on edge blocks
# after
from jax.experimental.mosaic.gpu import utils
col_c = min(col, width - 1) if isinstance(col, int) else col
x = tile[row, col_c]

shape = ir.MemRefType(ref.type).shape
assert all(-shape[a] <= i < shape[a] for a, i in enumerate(idx_tuple) if isinstance(i, int))

def in_bounds(shape, idx) -> bool:
    return all(-b <= i < b for b, i in zip(shape, idx) if isinstance(i, int))

try:
    v = tile[i, j]
except IndexError:
    v = zero  # predicated-off lane
    # or: continue

When it happens

Trigger: ref[idx] where idx is a Python int >= shape[axis] or < -shape[axis] with check_oob enabled; e.g. ref[64] on a shape (32, 8) memref, or ref[-33] where bound is 32.

Common situations: Using grid/block indices from launch geometry to index a shared-memory tile without clamping; shrinking a test tensor without updating indices; converting numpy prototype code that used out-of-bounds-but-never-hit indices.

Related errors

• Folding {fold_rank} dimensions starting from {dim} is out of
• Slice {idx} along axis {axis} is out of bounds for shape {sh
• SubViewOp only supports a single tile transform.
• Only support memref.cast where the input and output types ar
• memref.cast transforms must have identical transforms for bo

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