jax-ml/jax · error · IndexError

Out-of-bounds read of {allocation_key}: reading [{read_range

Error message

Out-of-bounds read of {allocation_key}: reading [{read_range}] but buffer has shape {shape}.

What it means

Raised by the Mosaic GPU interpreter when a Pallas kernel reads a block that extends beyond the bounds of an output/scratch buffer. The interpreter simulates each grid iteration on host memory, so any Block indexing that runs past the buffer shape is detected and rejected exactly as it would be an OOB access on device.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:473

    allocation_key: HostAllocationKey,
    read_range: tuple[int | slice, ...],
    shared_memory: memory.GPUSharedMemory,
    source_info,
    input_name: str | None,
    block_indices: tuple[int, ...] | None,
    grid_loop_idx: tuple[int, ...] | None,
) -> np.ndarray:
  """Handles out-of-bounds read based on shared_memory configuration."""
  if shared_memory.out_of_bounds_reads == "raise":
    if source_info is None:
      ctx = contextlib.nullcontext()
    else:
      ctx = source_info_util.user_context(
          traceback=source_info.traceback, name_stack=source_info.name_stack
      )
    with ctx:
      if input_name is None:
        raise IndexError(
            f"Out-of-bounds read of {allocation_key}:"
            f" reading [{read_range}] but buffer has shape {shape}."
        )
      else:
        # Different error message when we are reading a block of an input,
        # to copy it to a buffer before invoking the kernel body.
        raise IndexError(
            f"Out-of-bounds block index {block_indices} for {allocation_key},"
            f' input "{input_name}" in iteration {grid_loop_idx}:'
            f" reading [{read_range}] but input has shape {shape}."
        )
  # out_of_bounds_reads == "uninitialized"
  uninit_array = np.full(
      full_read_shape,
      interpret_utils.get_uninitialized_value(
          dtype, shared_memory.uninitialized_memory
      ),
      dtype=dtype,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the grid use ceil-div: grid = (shape + block - 1) // block so no iteration starts past the end
2. Add masked reads (e.g. use pallas masking / load with mask) or pad inputs so block divides shape exactly
3. Check index arithmetic in the kernel: ensure start_index + block_shape <= buffer.shape for the last iteration
4. Run under interpret mode while prototyping to catch OOB before hitting the device

Example fix

# before
grid = (n // bs,)  # misses remainder -> last block reads OOB if n % bs != 0
# after
grid = ((n + bs - 1) // bs,)  # ceil-div grid; kernel uses masked loads for tail

def check_bounds(shape, start, block):
    for s, st, b in zip(shape, start, block):
        assert st + b <= s, f'OOB: [{st}, {st+b}) vs dim {s}'
# before each interpret run, simulate:
for i in range(grid[0]):
    check_blocks(buffers, i, block_shapes)

try:
    out = kernel(x)
except IndexError as e:
    if 'Out-of-bounds read' in str(e):
        # fix grid/block sizing; print reported read_range vs shape
        ...

When it happens

Trigger: Calling a Mosaic GPU kernel under interpret mode (JAX_ENABLE_PERSISTENT_KERNEL_CACHE / pallas interpret or interpret mode activated) where a Block index computed from the grid iteration overruns the referenced buffer, e.g. block_start + block_size > dim for some iteration.

Common situations: Non-divisible grid: block sizes that don't evenly divide the tensor shape and the kernel lacks masked/padded reads; off-by-one in manual index arithmetic; using a grid larger than ceil(shape/block).

Related errors

• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• Range {rnge} is entirely out of bounds for shape {self.shape
• `thread_id` must be zero when allocating a buffer for all th
• `block_id` must be zero when allocating a buffer for all thr
• Out-of-bounds block index {block_indices} for {allocation_ke

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