jax-ml/jax · error · ValueError

Out-of-bounds masked swap of {allocation_key}: swapping [{re

Error message

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

What it means

Same out-of-bounds swap detection, but raised when a mask was provided — meaning some lanes where the mask is True still index outside the buffer, so even the masked swap cannot proceed.

Source

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

      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,
        read_write_range,
        source_info=source_info,
    )
  return token, ret


def call_swap(
    *,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the mask so it is False for every out-of-bounds index: mask = (idx < dim)
2. Check mask broadcasting/shape matches the Block's shape exactly
3. Verify you're comparing against the correct dimension of the target buffer
4. Add an interpret-mode unit test for boundary iterations

Example fix

# before
mask = idx < block_size  # wrong bound
# after
mask = idx < n  # n = actual buffer dim

import jax.numpy as jnp
idx = start + jnp.arange(bs)
mask = idx < n  # compare to buffer dim, not block size
assert not (mask & (idx >= n)).any()

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

When it happens

Trigger: A masked swap/atomic in interpret mode where the mask does not fully exclude out-of-bounds lanes (mask True for indices >= dim), e.g. mask built against the wrong dimension or wrong shape.

Common situations: Masks computed against block shape instead of buffer bounds; broadcast errors producing wrong-length masks; mixing up row/column masks in 2D kernels.

Related errors

• Out-of-bounds swap of {allocation_key}: swapping [{read_writ
• Out-of-bounds masked swap of ({device_id} {local_core_id} {m
• 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

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