{"record":{"id":"c89f150311724c55","repo":"jax-ml/jax","slug":"out-of-bounds-swap-of-device-id-local-core-id","errorCode":null,"errorMessage":"Out-of-bounds swap of ({device_id} {local_core_id} {memory_space} {buffer_id}): swapping [{read_write_range}] but buffer has shape {shape} .","messagePattern":"Out-of-bounds swap of \\((.+?) (.+?) (.+?) (.+?)\\): swapping \\[(.+?)\\] but buffer has shape (.+?) \\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":861,"sourceCode":"\n  key = (memory_space, buffer_id, device_id, local_core_id_for_buffer)\n  read_write_range = interpret_utils.to_range(transforms)\n  ret, (shape, _), clock = shared_memory.swap_buffer_content(\n      key,\n      read_write_range,\n      val,\n      mask,\n      (device_id, local_core_id),\n      logging_info=interpret_utils.TPULoggingInfo(\n          device_id=device_id,\n          local_core_id=local_core_id,\n          source_info=source_info,\n      ),\n  )\n\n  if ret is None:\n    if mask is None:\n      raise ValueError(\n          'Out-of-bounds swap of'\n          f' ({device_id} {local_core_id} {memory_space} {buffer_id}):'\n          f' swapping [{read_write_range}] but buffer has shape'\n          f' {shape} .'\n      )\n    else:\n      # TODO(jburnim): Include indices of out-of-bounds locations where mask\n      # is True.\n      raise ValueError(\n          'Out-of-bounds masked swap of'\n          f' ({device_id} {local_core_id} {memory_space} {buffer_id}): swapping'\n          f' [{read_write_range}] but buffer has shape {shape} . '\n      )\n\n  if shared_memory.detect_races:\n    assert races is not None\n    assert clock is not None\n    races.check_write(","sourceCodeStart":843,"sourceCodeEnd":879,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L843-L879","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nidx = i * BS  # grid one iteration too large\nold = sem_or_ref.swap(idx, val)\n# after\ngrid = (buf.shape[0] // BS,)  # ensure idx + BS <= buf.shape[0]","handlingStrategy":"validation","validationCode":"start = int(idx)\nassert 0 <= start and start + block_size <= buf_len, f'swap [{start},{start+block_size}) exceeds buffer {buf_len}'","typeGuard":null,"tryCatchPattern":"try:\n    old = ref.swap(idx, val)\nexcept ValueError as e:\n    if 'Out-of-bounds swap' in str(e):\n        # clamp or skip tail swap when partially out of range\n        raise","preventionTips":["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"],"tags":["jax","pallas","tpu","atomics","out-of-bounds"],"backgroundTag":"array-index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}