jax-ml/jax · error · ValueError

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

Error message

Out-of-bounds swap of ({device_id} {local_core_id} {memory_space} {buffer_id}): swapping [{read_write_range}] but buffer has shape {shape} .

What it means

The Pallas TPU interpreter performs swaps (atomic exchange) by reading and writing the referenced range; if the unmasked swap range [read_write_range] falls outside the buffer shape, it raises this ValueError because the simulated atomic exchange would touch out-of-bounds memory.

Source

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

  key = (memory_space, buffer_id, device_id, local_core_id_for_buffer)
  read_write_range = interpret_utils.to_range(transforms)
  ret, (shape, _), clock = shared_memory.swap_buffer_content(
      key,
      read_write_range,
      val,
      mask,
      (device_id, local_core_id),
      logging_info=interpret_utils.TPULoggingInfo(
          device_id=device_id,
          local_core_id=local_core_id,
          source_info=source_info,
      ),
  )

  if ret is None:
    if mask is None:
      raise ValueError(
          'Out-of-bounds swap of'
          f' ({device_id} {local_core_id} {memory_space} {buffer_id}):'
          f' swapping [{read_write_range}] but buffer has shape'
          f' {shape} .'
      )
    else:
      # TODO(jburnim): Include indices of out-of-bounds locations where mask
      # is True.
      raise ValueError(
          'Out-of-bounds masked swap of'
          f' ({device_id} {local_core_id} {memory_space} {buffer_id}): swapping'
          f' [{read_write_range}] but buffer has shape {shape} . '
      )

  if shared_memory.detect_races:
    assert races is not None
    assert clock is not None
    races.check_write(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Check that swap start index + block shape <= buffer shape on every grid iteration
Fix grid sizing relative to the buffer dimensions
Mask the swap (which yields the masked variant of the check) only if out-of-bounds lanes are intentional — otherwise correct indices

Example fix

# before
idx = i * BS  # grid one iteration too large
old = sem_or_ref.swap(idx, val)
# after
grid = (buf.shape[0] // BS,)  # ensure idx + BS <= buf.shape[0]

Defensive patterns

Strategy: validation

Validate before calling

start = int(idx)
assert 0 <= start and start + block_size <= buf_len, f'swap [{start},{start+block_size}) exceeds buffer {buf_len}'

Try / catch

try:
    old = ref.swap(idx, val)
except ValueError as e:
    if 'Out-of-bounds swap' in str(e):
        # clamp or skip tail swap when partially out of range
        raise

Prevention

Assert swap range fits buffer before the call
Size buffers to a multiple of the swapped block
Treat swap as read+write when computing bounds

When it happens

Trigger: An atomic swap on a buffer where the computed index range exceeds the buffer's shape — e.g., swap indices + swapped block size beyond the buffer extent, with no mask limiting the accessed region.

Common situations: Grid/index math errors in kernels using semaphores or atomics; buffers sized smaller than the swapped block; forgetting that swap reads AND writes the full range.

Related errors

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