jax-ml/jax · error · IndexError

Out-of-bounds block index {block_indices} for output "{outpu

Error message

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

What it means

While copying a kernel buffer's result back into a block of a pallas_call output after kernel invocation, the interpreter found the block index for that output iteration reads outside the output's shape. It names the output, grid iteration, device, and core where the violation occurred.

Source

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

      logging_info=interpret_utils.TPULoggingInfo(
          device_id=device_id,
          local_core_id=local_core_id,
          source_info=source_info,
      ),
  )
  clock = clock if clock is not None else clock_

  if not in_bounds:
    if output_name is None:
      raise ValueError(
          'Out-of-bounds write of'
          f' ({device_id} {local_core_id} {memory_space} {buffer_id}):'
          f' writing [{write_range}] but buffer has shape {shape} .'
      )
    else:
      # Different error message when we are copying a kernel buffer to a
      # block of an output (just after a kernel invocation).
      raise IndexError(
          f'Out-of-bounds block index {block_indices} for'
          f' output "{output_name}" in iteration {grid_loop_idx}'
          f' on device {device_id} (core {local_core_id}):'
          f' reading [{write_range}] but output has shape {shape}.'
      )

  if shared_memory.detect_races:
    if src_device_id is None:
      src_device_id = device_id
    if src_local_core_id is None:
      src_local_core_id = local_core_id
    assert races is not None
    assert clock is not None
    races.check_write(
        (src_device_id, src_local_core_id),
        clock,
        (memory_space, buffer_id, device_id, local_core_id_for_buffer),
        write_range,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Align the grid with the output's block count (ceil(out.shape / block_shape))
Pad outputs to multiples of the block shape
Fix the output index_map so start + block_shape <= out.shape on every iteration

Example fix

# before
out_spec = BlockSpec((BN,), lambda i: (i * BN,))
grid = (N_total // BM,)  # sized for input, too big for output
# after
grid = (min(x.shape[0], out.shape[0]) // BM,)  # or pad both to a common multiple

Defensive patterns

Strategy: validation

Validate before calling

import math
n_out_blocks = math.ceil(out.shape[0] / out_spec.block_shape[0])
assert grid[0] <= n_out_blocks, f'grid {grid} > output blocks {n_out_blocks}'

Try / catch

try:
    run(inputs)
except IndexError as e:
    if 'Out-of-bounds block index' in str(e) and 'output' in str(e):
        pad_output_and_retry()

Prevention

Derive grid from output as well as input block counts
Keep input/output shapes consistent or padded to common block multiples
Verify output index_map in interpret mode before hardware runs

When it happens

Trigger: An output BlockSpec whose block shape/index_map addresses beyond out.shape on some grid iteration — typically a grid larger than the number of output blocks, or a non-divisible output size without padding.

Common situations: Different input/output sizes with one shared grid; output tensors smaller than the block shape; index_map bugs scaling program_id too far for the output.

Related errors

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