jax-ml/jax · error · IndexError

Range {rnge} is entirely out of bounds for shape {self.shape

Error message

Range {rnge} is entirely out of bounds for shape {self.shape}.

What it means

Interpret-mode Buffer.__getitem__ received an index range that lies entirely outside the allocated shape. Since even the clipped intersection is empty (clip returned None), reading is meaningless and the error is raised rather than returning uninitialized data.

Source

Thrown at jax/_src/pallas/mosaic/interpret/shared_memory.py:317

    """Returns the portion of `self.content` specified by `rnge`.

    Args:
      rnge: The range to read.

    Raises:
      IndexError: If `rnge` is entirely out of bounds for the allocated array in
        the `Buffer`, i.e. `rnge` is out of bounds for `self.shape`.

    Returns:
      The portion of `self.content` specified by `rnge`.
    """
    rnge = self._normalize_range(rnge)
    rnge_or_none = interpret_utils.clip_range_to_shape(rnge, self.shape)
    if rnge_or_none is None:
      # Raise if reading entirely outside of the allocated shape. We leave it to
      # the client to handle the case where out-of-bounds reads are allowed (and
      # should return uninitialized values).
      raise IndexError(
          f"Range {rnge} is entirely out of bounds for shape {self.shape}."
      )

    rnge = rnge_or_none
    return self.content[rnge]

  def _set_within_logical_shape(
      self, rnge: tuple[slice | int, ...], value: np.ndarray
  ):
    """Updates `self.content` with `value` for the portion of `rnge` within `self.logical_shape`."""
    rnge_or_none = interpret_utils.clip_range_to_shape(rnge, self.logical_shape)
    if rnge_or_none is None:
      return

    rnge = rnge_or_none
    shape_to_write = self.content[rnge].shape
    self.content[rnge] = value[tuple(slice(0, s) for s in shape_to_write)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Fix the index computation so slice starts fall within the allocation (clamp with min(start, size - 1))
  2. Verify grid iteration bounds match the allocation shape; shrink grid or enlarge allocation
  3. If OOB reads are expected to return uninitialized values, route reads through APIs that handle padding (_get_with_padding / out-of-bounds handling)

Example fix

// before
val = smem[start:block_size, :]
// after
val = smem[min(start, smem.shape[0] - 1):min(start + block_size, smem.shape[0]), :]
Defensive patterns

Strategy: validation

Validate before calling

from jax._src.pallas.mosaic.interpret import interpret_utils
rnge = tuple(r if isinstance(r, slice) else slice(r, r + 1) for r in rnge)
assert interpret_utils.clip_range_to_shape(rnge, buf.shape) is not None, 'read range entirely OOB'

Try / catch

try:
    data = buf[rnge]
except IndexError as e:
    if 'entirely out of bounds' in str(e):
        data = np.zeros(())  # or skip tile
    else:
        raise

Prevention

When it happens

Trigger: Indexing an interpret shared memory Buffer with a slice whose start is >= the allocation extent on some axis, e.g. buf[128:256, :] on a shape-(128, 64) buffer, when no out-of-bounds-read permission was arranged by the caller.

Common situations: Grid/block-index arithmetic that walks past the end of an allocation; incorrect grid size assumptions on edge tiles; kernels relying on padded OOB reads without configuring interpret mode to allow them.

Related errors


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