jax-ml/jax · error · IndexError

Out-of-bounds block index {block_indices} for {allocation_ke

Error message

Out-of-bounds block index {block_indices} for {allocation_key}, input "{input_name}" in iteration {grid_loop_idx}: reading [{read_range}] but input has shape {shape}.

What it means

Same out-of-bounds detection as the plain read case, but specialized to input blocks that the interpreter must copy into a buffer before running the kernel body. It reports the block indices, grid loop iteration, requested read range, and the actual input shape.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use masked loads/stores (Block with mask or swap in bounded index computation) for boundary iterations
Pad the input to a multiple of the block size before passing it to the kernel
Verify the Block's dimension maps produce indices within [0, dim) for every grid iteration
Reproduce in interpret mode with tiny shapes to identify the offending iteration printed in the message

Example fix

# before
x_block = x[ds[start, bs], :]  # unmasked tail read
# after
mask = (start + jnp.arange(bs)) < n
x_block = x[ds[start, bs], :].mask(mask, 0.0)  # or pad x to multiple of bs

Defensive patterns

Strategy: validation

Validate before calling

n = x.shape[0]
assert all((i*bs + bs <= n) or use_mask for i in range(nblocks))

Try / catch

try:
    out = kernel(x)
except IndexError as e:
    if 'Out-of-bounds block index' in str(e):
        pad or mask inputs

Prevention

Pad inputs to multiples of block sizes
Default to masked loads on boundary blocks

When it happens

Trigger: A kernel in interpret mode indexes an input Block (via a transformed reference) such that for some grid iteration the read range exceeds the input's shape; typical when grid is ceil-div but reads are unmasked.

Common situations: Tail iterations of non-evenly-divisible shapes; kernels written assuming padded inputs (e.g. from a pipeline with padded batch dim) run without padding; incorrect Block reference constructed from dmapped user indices.

Related errors

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