jax-ml/jax · error · ValueError

Out-of-bounds masked swap of

Error message

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

What it means

Even with a mask, a Pallas TPU interpreter swap must not have any out-of-bounds location where the mask is True: the interpreter validates that the mask never enables a lane whose address lies outside the buffer. If the overall range is out of bounds and the mask could touch those lanes, this ValueError fires.

Solutions

Make the mask index-based: mask = (idx + arange(B) < n), not value-based
Verify mask shape equals the swapped value shape and correctly broadcasts
Ensure mask is False for all lanes whose addresses fall outside the buffer

Example fix

# before
mask = vals != 0  # may be True out of bounds
old = ref.swap(idx, vals, mask=mask)
# after
mask = (idx + jnp.arange(B)) < n
old = ref.swap(idx, vals, mask=mask)

Defensive patterns

Strategy: type-guard

Validate before calling

rows = idx + jnp.arange(block_size)
mask = rows < n  # index-based mask, guaranteed False out of bounds
assert not jnp.any(mask & (rows >= buf_len))

Type guard

def mask_is_index_based(mask, idx, n) -> bool:
    return bool(jnp.all((idx + jnp.arange(mask.shape[-1]))[..., None, :] .squeeze(-2) is not None)) or bool(jnp.all((idx + jnp.arange(mask.shape[-1])) < n) | ~mask.any())

Try / catch

try:
    old = ref.swap(idx, val, mask=mask)
except ValueError as e:
    if 'Out-of-bounds masked swap' in str(e):
        mask = (idx + jnp.arange(B)) < n
        old = ref.swap(idx, val, mask=mask)

Prevention

Always derive masks from index bounds, not data values
Zero/pad buffers so out-of-bounds lanes never matter
Test masked swaps at grid boundaries in interpret mode

When it happens

Trigger: A masked swap (e.g., boundary handling in an atomic update) where the mask array is True at positions outside the buffer — commonly a mask computed from a value pattern rather than from index bounds.

Common situations: Boundary handling masks like (vals != 0) instead of (row < n); mask shape/broadcast misalignment causing True in out-of-bounds lanes; forgetting that padded lanes must be masked off by index.

Related errors

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

Appendix: source

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

      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(
        (device_id, local_core_id),
        clock,
        (memory_space, buffer_id, device_id, local_core_id_for_buffer),
        read_write_range,
        source_info=source_info,
    )
  return token, ret

View on GitHub (pinned to 1e1c6a8fc0)