{"record":{"id":"cc2739e0d5054720","repo":"jax-ml/jax","slug":"out-of-bounds-masked-swap-of-allocation-key-swa","errorCode":null,"errorMessage":"Out-of-bounds masked swap of {allocation_key}: swapping [{read_write_range}] but buffer has shape {shape} . ","messagePattern":"Out-of-bounds masked swap of (.+?): swapping \\[(.+?)\\] but buffer has shape (.+?) \\. ","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py","lineNumber":685,"sourceCode":"      np.array(val),\n      np.array(mask) if mask is not None else None,\n      thread,\n      increment_clock=increment_clock,\n      logging_info=memory.GPULoggingInfo(mesh_location, thread, source_info),\n  )\n  clock = clock if clock is not None else clock_\n\n  if ret is None:\n    if mask is None:\n      raise ValueError(\n          f\"Out-of-bounds swap of {allocation_key}:\"\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          f\"Out-of-bounds masked swap of {allocation_key}: swapping\"\n          f\" [{read_write_range}] but buffer has shape {shape} . \"\n      )\n\n  if shared_memory.detect_races:\n    assert clock is not None\n    get_races().check_write(\n        thread,\n        clock.generic_clock,\n        allocation_key,\n        read_write_range,\n        source_info=source_info,\n    )\n  return token, ret\n\n\ndef call_swap(\n    *,","sourceCodeStart":667,"sourceCodeEnd":703,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py#L667-L703","documentation":"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.","triggerScenarios":"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.","commonSituations":"Masks computed against block shape instead of buffer bounds; broadcast errors producing wrong-length masks; mixing up row/column masks in 2D kernels.","solutions":["Fix the mask so it is False for every out-of-bounds index: mask = (idx < dim)","Check mask broadcasting/shape matches the Block's shape exactly","Verify you're comparing against the correct dimension of the target buffer","Add an interpret-mode unit test for boundary iterations"],"exampleFix":"# before\nmask = idx < block_size  # wrong bound\n# after\nmask = idx < n  # n = actual buffer dim","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nidx = start + jnp.arange(bs)\nmask = idx < n  # compare to buffer dim, not block size\nassert not (mask & (idx >= n)).any()","typeGuard":null,"tryCatchPattern":"try:\n    kernel(x)\nexcept ValueError as e:\n    if 'Out-of-bounds masked swap' in str(e):\n        fix mask bounds","preventionTips":["Always build masks against the true dimension","Unit-test masks at boundary iterations"],"tags":["pallas","mosaic-gpu","interpret-mode","out-of-bounds","atomics","mask"],"backgroundTag":"buffer-out-of-bounds-write","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}