{"record":{"id":"aa0c28947896e747","repo":"jax-ml/jax","slug":"out-of-bounds-masked-swap-of-device-id-local-c","errorCode":null,"errorMessage":"Out-of-bounds masked swap of ({device_id} {local_core_id} {memory_space} {buffer_id}): swapping [{read_write_range}] but buffer has shape {shape} . ","messagePattern":"Out-of-bounds masked 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":870,"sourceCode":"      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(\n        (device_id, local_core_id),\n        clock,\n        (memory_space, buffer_id, device_id, local_core_id_for_buffer),\n        read_write_range,\n        source_info=source_info,\n    )\n  return token, ret\n\n","sourceCodeStart":852,"sourceCodeEnd":888,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L852-L888","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nmask = vals != 0  # may be True out of bounds\nold = ref.swap(idx, vals, mask=mask)\n# after\nmask = (idx + jnp.arange(B)) < n\nold = ref.swap(idx, vals, mask=mask)","handlingStrategy":"type-guard","validationCode":"rows = idx + jnp.arange(block_size)\nmask = rows < n  # index-based mask, guaranteed False out of bounds\nassert not jnp.any(mask & (rows >= buf_len))","typeGuard":"def mask_is_index_based(mask, idx, n) -> bool:\n    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())","tryCatchPattern":"try:\n    old = ref.swap(idx, val, mask=mask)\nexcept ValueError as e:\n    if 'Out-of-bounds masked swap' in str(e):\n        mask = (idx + jnp.arange(B)) < n\n        old = ref.swap(idx, val, mask=mask)","preventionTips":["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"],"tags":["jax","pallas","tpu","atomics","mask","out-of-bounds"],"backgroundTag":"invalid-mask","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}