jax-ml/jax · error · ValueError

Out-of-bounds swap of {allocation_key}: swapping [{read_writ

Error message

Out-of-bounds swap of {allocation_key}: swapping [{read_write_range}] but buffer has shape {shape} .

What it means

The interpreter's swap (atomic read-modify-write) callback detected that the range being swapped falls outside the buffer shape and no mask was provided, so the access cannot be made safe.

Source

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

    assert mask.shape == val.shape

  shared_memory = _get_shared_memory()

  read_write_range = interpret_utils.to_range(transforms)
  ret, (shape, _), clock_ = shared_memory.swap_buffer_content(
      allocation_key,
      read_write_range,
      np.array(val),
      np.array(mask) if mask is not None else None,
      thread,
      increment_clock=increment_clock,
      logging_info=memory.GPULoggingInfo(mesh_location, thread, source_info),
  )
  clock = clock if clock is not None else clock_

  if ret is None:
    if mask is None:
      raise ValueError(
          f"Out-of-bounds swap of {allocation_key}:"
          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(
          f"Out-of-bounds masked swap of {allocation_key}: swapping"
          f" [{read_write_range}] but buffer has shape {shape} . "
      )

  if shared_memory.detect_races:
    assert clock is not None
    get_races().check_write(
        thread,
        clock.generic_clock,
        allocation_key,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add a mask to the swap/atomic operation so boundary lanes are disabled
2. Fix the grid to ceil-div and ensure the swap index arithmetic stays in bounds
3. Pad the target buffer to a multiple of the block shape
4. Reproduce with small shapes in interpret mode to find the offending iteration

Example fix

# before
dst[ds[start, bs]].atomic_add(val)  # unmasked
# after
mask = (start + jnp.arange(bs)) < n
dst[ds[start, bs]].atomic_add(val, mask=mask)

assert all(start + bs <= dim for start, dim in zip(starts, dst.shape))

try:
    kernel(x)
except ValueError as e:
    if 'Out-of-bounds swap' in str(e):
        add mask or fix indices

When it happens

Trigger: Calling a swap/atomic op (e.g. atomic_add on a Block) in interpret mode where the index range exceeds the buffer; specifically the unmasked branch (mask is None).

Common situations: Boundary iterations with non-divisible shapes and unmasked atomics; wrong grid size causing out-of-range swaps on accumulator buffers.

Related errors

• Out-of-bounds masked swap of {allocation_key}: swapping [{re
• Out-of-bounds read of {allocation_key}: reading [{read_range
• Out-of-bounds block index {block_indices} for {allocation_ke
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• Out-of-bounds swap of ({device_id} {local_core_id} {memory_s

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