{"record":{"id":"1179fa6b3d347744","repo":"jax-ml/jax","slug":"out-of-bounds-read-of-device-id-local-core-id","errorCode":null,"errorMessage":"Out-of-bounds read of ({device_id} {local_core_id} {memory_space} {buffer_id}): reading [{read_range}] but buffer has shape {shape}.","messagePattern":"Out-of-bounds read of \\((.+?) (.+?) (.+?) (.+?)\\): reading \\[(.+?)\\] but buffer has shape (.+?)\\.","errorType":"validation","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":673,"sourceCode":"      full_read_shape.append(dim_size)\n    elif isinstance(idx_or_slice, int):\n      continue\n    else:\n      dim_size = (idx_or_slice.stop - idx_or_slice.start) // idx_or_slice.step\n      assert isinstance(dim_size, int)\n      full_read_shape.append(dim_size)\n\n  if (ret is None) or (tuple(full_read_shape) != ret.shape):\n    if shared_memory.out_of_bounds_reads == 'raise':\n      if source_info is None:\n        ctx = contextlib.nullcontext()\n      else:\n        ctx = source_info_util.user_context(\n            traceback=source_info.traceback, name_stack=source_info.name_stack\n        )\n      with ctx:\n        if input_name is None:\n          raise IndexError(\n              'Out-of-bounds read of'\n              f' ({device_id} {local_core_id} {memory_space} {buffer_id}):'\n              f' reading [{read_range}] but buffer has shape {shape}.'\n          )\n        else:\n          # Different error message when we are reading a block of an input,\n          # to copy it to a buffer before invoking the kernel body.\n          raise IndexError(\n              f'Out-of-bounds block index {block_indices} for'\n              f' input \"{input_name}\" in iteration {grid_loop_idx}'\n              f' on device {device_id} (core {local_core_id}):'\n              f' reading [{read_range}] but input has shape {shape}.'\n          )\n    # out_of_bounds_reads == \"uninitialized\"\n    uninit_array = np.full(\n        full_read_shape,\n        interpret_utils.get_uninitialized_value(\n            dtype, shared_memory.uninitialized_memory","sourceCodeStart":655,"sourceCodeEnd":691,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L655-L691","documentation":"Raised by the TPU Pallas interpret-mode simulator when a kernel reads a VMEM/SMEM buffer slice outside its bounds. The interpreter emulates Pallas kernels on CPU/Numpy and checks every load's index range against the allocated buffer shape; an out-of-range read (e.g., from block indices exceeding the grid implied by the BlockSpec) triggers this IndexError with device/core/buffer details.","triggerScenarios":"A pallas_call kernel whose BlockSpec block shape or start indices computed per grid iteration address past the end of a reference buffer; e.g., grid larger than num_blocks derived from input shape, or a manual start-index computation like pl.program_id(0)*BM + offset beyond the buffer.","commonSituations":"Mismatch between grid size and input shape (e.g., input not divisible by block size so an extra iteration overreads), off-by-one in block index math, or incorrect BlockSpec index_map returning indices past valid blocks.","solutions":["Check that grid = ceil(div shapes / block shapes) matches the BlockSpec index_map so indices never exceed (shape - block_shape)/block_step","Run under interpret mode with small shapes to reproduce and print block_indices/start indices per iteration","Pad inputs to a multiple of the block size or adjust the grid to avoid the trailing partial block","Fix off-by-one errors in manual start-index arithmetic inside the kernel"],"exampleFix":"# before\ngrid = (input.shape[0] // BM + 1,)  # overreads on last block\n# after\ngrid = (input.shape[0] // BM,)  # or pad input to multiple of BM","handlingStrategy":"validation","validationCode":"import math\nassert all(math.ceil(d.shape[i] / bs[i]) >= grid[i] for spec, d in zip(in_specs, inputs) for i, bs in [enumerate_blocks(spec)]) or True\n# simpler: assert block fits\nfor x, spec in zip(inputs, in_specs):\n    for dim, blk in zip(x.shape, spec.block_shape):\n        assert blk <= dim, f'block {blk} > dim {dim}'","typeGuard":null,"tryCatchPattern":"try:\n    out = f(x)  # interpret-mode pallas_call\nexcept IndexError as e:\n    if 'Out-of-bounds read' in str(e):\n        # log grid/block shapes and shrink grid or pad inputs\n        raise","preventionTips":["Compute grid as ceil_div(shape, block) consistently for all inputs","Pad inputs to multiples of block shapes before pallas_call","Unit-test kernels in interpret mode with tiny shapes to catch index bugs cheaply"],"tags":["jax","pallas","tpu","out-of-bounds","interpret-mode"],"backgroundTag":"array-index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}