{"record":{"id":"f7953d4d26fde452","repo":"jax-ml/jax","slug":"out-of-bounds-write-of-device-id-local-core-id","errorCode":null,"errorMessage":"Out-of-bounds write of ({device_id} {local_core_id} {memory_space} {buffer_id}): writing [{write_range}] but buffer has shape {shape} .","messagePattern":"Out-of-bounds write of \\((.+?) (.+?) (.+?) (.+?)\\): writing \\[(.+?)\\] but buffer has shape (.+?) \\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":780,"sourceCode":"\n  key = (memory_space, buffer_id, device_id, local_core_id_for_buffer)\n  write_range = interpret_utils.to_range(transforms)\n  in_bounds, (shape, _), clock_ = shared_memory.store_buffer_content(\n      key,\n      write_range,\n      val,\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  clock = clock if clock is not None else clock_\n\n  if not in_bounds:\n    if output_name is None:\n      raise ValueError(\n          'Out-of-bounds write of'\n          f' ({device_id} {local_core_id} {memory_space} {buffer_id}):'\n          f' writing [{write_range}] but buffer has shape {shape} .'\n      )\n    else:\n      # Different error message when we are copying a kernel buffer to a\n      # block of an output (just after a kernel invocation).\n      raise IndexError(\n          f'Out-of-bounds block index {block_indices} for'\n          f' output \"{output_name}\" in iteration {grid_loop_idx}'\n          f' on device {device_id} (core {local_core_id}):'\n          f' reading [{write_range}] but output has shape {shape}.'\n      )\n\n  if shared_memory.detect_races:\n    if src_device_id is None:\n      src_device_id = device_id\n    if src_local_core_id is None:","sourceCodeStart":762,"sourceCodeEnd":798,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L762-L798","documentation":"The Pallas TPU interpreter detected a store (write) whose destination range falls outside the allocated buffer: writing [write_range] into a buffer of the given shape. This is the write-side counterpart of the out-of-bounds read check and fires when block indices/offsets plus block shape exceed the buffer dimensions.","triggerScenarios":"A pallas kernel store operation where computed start indices + block shape exceed the output buffer shape — e.g., grid sized larger than the output supports, or manual index math writing past the last valid block.","commonSituations":"Output shape not divisible by block size with a grid rounding up; off-by-one in the output index computation; mismatched BlockSpec for outputs versus actual out_shapes passed to pallas_call.","solutions":["Match the grid to the output's valid block count (ceil or floor consistently with your BlockSpec)","Verify output BlockSpec block shapes exactly equal what the kernel writes","Pad the output or adjust index_map so writes stay in bounds","Debug in interpret mode printing store indices per iteration"],"exampleFix":"# before\ngrid = (out.shape[0] // BN + 1,)\n# after\ngrid = (out.shape[0] // BN,)  # or pad out.shape[0] to a multiple of BN","handlingStrategy":"validation","validationCode":"import math\nfor out, spec in zip(out_shapes, out_specs):\n    for dim, blk in zip(out.shape, spec.block_shape):\n        assert blk <= dim\nassert all(math.ceil(o.shape[0]/s.block_shape[0]) >= grid[0] for o, s in zip(out_shapes, out_specs))","typeGuard":null,"tryCatchPattern":"try:\n    pallas_call(kernel, grid, out_shapes, inputs, interpret=True)(inputs)\nexcept ValueError as e:\n    if 'Out-of-bounds write' in str(e):\n        # recheck output BlockSpec and grid, pad output, retry\n        raise","preventionTips":["Match output grid to ceil(out.shape / out_block)","Ensure output block shapes equal what the kernel writes","Prefer padding over boundary-exceeding grids"],"tags":["jax","pallas","tpu","out-of-bounds","write"],"backgroundTag":"array-index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}