jax-ml/jax · error · IndexError

Out-of-bounds block index {block_indices} for input "{input_

Error message

Out-of-bounds block index {block_indices} for input "{input_name}" in iteration {grid_loop_idx} on device {device_id} (core {local_core_id}): reading [{read_range}] but input has shape {shape}.

What it means

The Pallas TPU interpreter detected that, when copying a block of a pallas_call input into a kernel buffer before invoking the kernel body, the block's read range (start index + block shape) exceeds the input's actual shape. It reports which input, grid iteration, device, and core produced the violation.

Source

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

  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
        ),
        dtype=dtype,
    )
    if ret is None:
      ret = uninit_array
    else:
      uninit_array[tuple(slice(s) for s in ret.shape)] = ret
      ret = uninit_array

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure every input's BlockSpec yields exactly ceil(input.shape / block_shape) valid blocks matching the grid
  2. Pad smaller inputs so their block counts align with the grid
  3. Fix the index_map to clamp or correctly compute per-input block offsets
  4. Reproduce in interpret mode and inspect the reported block_indices and read_range vs the input shape

Example fix

# before
in_spec = BlockSpec((BM,), lambda i: (i * BM,))  # grid too big for this input
# after
in_spec = BlockSpec((BM,), lambda i: (i * BM,))
grid = (x.shape[0] // BM,)  # match grid to smallest valid block count (pad x if needed)
Defensive patterns

Strategy: validation

Validate before calling

import math
def check_blockspec(x, spec, grid):
    n_blocks = math.ceil(x.shape[0] / spec.block_shape[0])
    assert grid[0] <= n_blocks, f'grid {grid} exceeds {n_blocks} blocks for input shape {x.shape}'

Try / catch

try:
    kernel_grid_run(inputs)
except IndexError as e:
    if 'Out-of-bounds block index' in str(e) and 'input' in str(e):
        pad_and_retry()

Prevention

When it happens

Trigger: A pallas_call input with a BlockSpec whose block shape or index_map produces start indices such that start + block_shape > input.shape on some grid iteration; typical when the grid is larger than the number of valid blocks for that input.

Common situations: Multiple inputs with different sizes but a grid sized for the largest; non-divisible shapes without padding; index_map that scales program_id beyond the valid block count for one specific input.

Related errors

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