jax-ml/jax · error · IndexError

Out-of-bounds read of ({device_id} {local_core_id} {memory_s

Error message

Out-of-bounds read of ({device_id} {local_core_id} {memory_space} {buffer_id}): reading [{read_range}] but buffer has shape {shape}.

What it means

Raised by the TPU Pallas interpret-mode simulator when a kernel reads a VMEM/SMEM buffer slice outside its bounds. The interpreter emulates Pallas kernels on CPU/Numpy and checks every load's index range against the allocated buffer shape; an out-of-range read (e.g., from block indices exceeding the grid implied by the BlockSpec) triggers this IndexError with device/core/buffer details.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:673

      full_read_shape.append(dim_size)
    elif isinstance(idx_or_slice, int):
      continue
    else:
      dim_size = (idx_or_slice.stop - idx_or_slice.start) // idx_or_slice.step
      assert isinstance(dim_size, int)
      full_read_shape.append(dim_size)

  if (ret is None) or (tuple(full_read_shape) != ret.shape):
    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(
              'Out-of-bounds read of'
              f' ({device_id} {local_core_id} {memory_space} {buffer_id}):'
              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'
              f' input "{input_name}" in iteration {grid_loop_idx}'
              f' on device {device_id} (core {local_core_id}):'
              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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Check that grid = ceil(div shapes / block shapes) matches the BlockSpec index_map so indices never exceed (shape - block_shape)/block_step
  2. Run under interpret mode with small shapes to reproduce and print block_indices/start indices per iteration
  3. Pad inputs to a multiple of the block size or adjust the grid to avoid the trailing partial block
  4. Fix off-by-one errors in manual start-index arithmetic inside the kernel

Example fix

# before
grid = (input.shape[0] // BM + 1,)  # overreads on last block
# after
grid = (input.shape[0] // BM,)  # or pad input to multiple of BM
Defensive patterns

Strategy: validation

Validate before calling

import math
assert all(math.ceil(d.shape[i] / bs[i]) >= grid[i] for spec, d in zip(in_specs, inputs) for i, bs in [enumerate_blocks(spec)]) or True
# simpler: assert block fits
for x, spec in zip(inputs, in_specs):
    for dim, blk in zip(x.shape, spec.block_shape):
        assert blk <= dim, f'block {blk} > dim {dim}'

Try / catch

try:
    out = f(x)  # interpret-mode pallas_call
except IndexError as e:
    if 'Out-of-bounds read' in str(e):
        # log grid/block shapes and shrink grid or pad inputs
        raise

Prevention

When it happens

Trigger: A pallas_call kernel whose BlockSpec block shape or start indices computed per grid iteration address past the end of a reference buffer; e.g., grid larger than num_blocks derived from input shape, or a manual start-index computation like pl.program_id(0)*BM + offset beyond the buffer.

Common situations: Mismatch between grid size and input shape (e.g., input not divisible by block size so an extra iteration overreads), off-by-one in block index math, or incorrect BlockSpec index_map returning indices past valid blocks.

Related errors

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