jax-ml/jax · error · IndexError

Range {rnge} is (at least partially) out of bounds for alloc

Error message

Range {rnge} is (at least partially) out of bounds for allocation shape {self.shape}.

What it means

Interpret-mode Buffer.__setitem__ received a write range that is at least partially outside the allocation shape. Unlike reads, any partial out-of-bounds write is rejected because it would silently corrupt or drop data.

Source

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

    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)]

  def __setitem__(self, rnge: tuple[slice | int, ...], value: np.ndarray):
    """Updates `self.content` with `value`, if `rnge` is fully within `self.shape`.

    Args:
      rnge: The range to write.
      value: The value to write.

    Raises:
      IndexError: If any part of `rnge` is out of bounds for the allocated array
        in the `Buffer`, i.e. if any part of `rnge` is out of bounds for
        `self.shape`.
    """
    rnge = self._normalize_range(rnge)
    if interpret_utils.is_range_out_of_bounds_for_shape(rnge, self.shape):
      raise IndexError(
          f"Range {rnge} is (at least partially) out of bounds for"
          f" allocation shape {self.shape}."
      )

    self._set_within_logical_shape(rnge, value)

  def set_in_bounds_portion(
      self, rnge: tuple[slice | int, ...], value: np.ndarray
  ):
    """Updates `self.content` with `value` for the portion of `rnge` within `self.logical_shape`."""
    rnge = self._normalize_range(rnge)
    self._set_within_logical_shape(rnge, value)


@dataclasses.dataclass(frozen=True)
class ShapeAndDtype:
  shape: Sequence[int]
  dtype: np.dtype

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clip the write range to the allocation before storing (slice(min(start, size), min(stop, size)))
2. Add proper masking so partial edge tiles only write in-bounds elements
3. Enlarge the allocation or fix the block/grid size math so ranges fit

Example fix

// before
smem[off:off + block_size] = value
// after
smem[off:min(off + block_size, smem.shape[0])] = value[:min(block_size, smem.shape[0] - off)]

from jax._src.pallas.mosaic.interpret import interpret_utils
assert not interpret_utils.is_range_out_of_bounds_for_shape(rnge, buf.shape), 'write range partially OOB'

When it happens

Trigger: Writing to an interpret shared memory Buffer with a slice extending past the allocation extent, e.g. smem[112:128] on a 120-element buffer.

Common situations: Boundary tiles in tiled kernels where (index * block_size + block_size) exceeds the allocation; mismatched grid and buffer sizes; incorrect masking in store operations.

Related errors

• Range {rnge} is entirely out of bounds for shape {self.shape
• Out-of-bounds write of ({device_id} {local_core_id} {memory_
• Acc ref must be at least 2D, got shape {shape}
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• Out-of-bounds block index {block_indices} for input "{input_

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