jax-ml/jax · error · ValueError

Out-of-bounds write of ({device_id} {local_core_id} {memory_

Error message

Out-of-bounds write of ({device_id} {local_core_id} {memory_space} {buffer_id}): writing [{write_range}] but buffer has shape {shape} .

What it means

The Pallas TPU interpreter detected a store (write) whose destination range falls outside the allocated buffer: writing [write_range] into a buffer of the given shape. This is the write-side counterpart of the out-of-bounds read check and fires when block indices/offsets plus block shape exceed the buffer dimensions.

Source

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

  key = (memory_space, buffer_id, device_id, local_core_id_for_buffer)
  write_range = interpret_utils.to_range(transforms)
  in_bounds, (shape, _), clock_ = shared_memory.store_buffer_content(
      key,
      write_range,
      val,
      (device_id, local_core_id),
      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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Match the grid to the output's valid block count (ceil or floor consistently with your BlockSpec)
Verify output BlockSpec block shapes exactly equal what the kernel writes
Pad the output or adjust index_map so writes stay in bounds
Debug in interpret mode printing store indices per iteration

Example fix

# before
grid = (out.shape[0] // BN + 1,)
# after
grid = (out.shape[0] // BN,)  # or pad out.shape[0] to a multiple of BN

Defensive patterns

Strategy: validation

Validate before calling

import math
for out, spec in zip(out_shapes, out_specs):
    for dim, blk in zip(out.shape, spec.block_shape):
        assert blk <= dim
assert all(math.ceil(o.shape[0]/s.block_shape[0]) >= grid[0] for o, s in zip(out_shapes, out_specs))

Try / catch

try:
    pallas_call(kernel, grid, out_shapes, inputs, interpret=True)(inputs)
except ValueError as e:
    if 'Out-of-bounds write' in str(e):
        # recheck output BlockSpec and grid, pad output, retry
        raise

Prevention

Match output grid to ceil(out.shape / out_block)
Ensure output block shapes equal what the kernel writes
Prefer padding over boundary-exceeding grids

When it happens

Trigger: A pallas kernel store operation where computed start indices + block shape exceed the output buffer shape — e.g., grid sized larger than the output supports, or manual index math writing past the last valid block.

Common situations: Output shape not divisible by block size with a grid rounding up; off-by-one in the output index computation; mismatched BlockSpec for outputs versus actual out_shapes passed to pallas_call.

Related errors

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