{"record":{"id":"4cff24ba83a35528","repo":"jax-ml/jax","slug":"out-of-bounds-block-index-block-indices-for-inpu","errorCode":null,"errorMessage":"Out-of-bounds block index {block_indices} for input \"{input_name}\" in iteration {grid_loop_idx} on device {device_id} (core {local_core_id}): reading [{read_range}] but input has shape {shape}.","messagePattern":"Out-of-bounds block index (.+?) for input \"(.+?)\" in iteration (.+?) on device (.+?) \\(core (.+?)\\): reading \\[(.+?)\\] but input has shape (.+?)\\.","errorType":"validation","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":681,"sourceCode":"  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\n        ),\n        dtype=dtype,\n    )\n    if ret is None:\n      ret = uninit_array\n    else:\n      uninit_array[tuple(slice(s) for s in ret.shape)] = ret\n      ret = uninit_array","sourceCodeStart":663,"sourceCodeEnd":699,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L663-L699","documentation":"The Pallas TPU interpreter detected that, when copying a block of a pallas_call input into a kernel buffer before invoking the kernel body, the block's read range (start index + block shape) exceeds the input's actual shape. It reports which input, grid iteration, device, and core produced the violation.","triggerScenarios":"A pallas_call input with a BlockSpec whose block shape or index_map produces start indices such that start + block_shape > input.shape on some grid iteration; typical when the grid is larger than the number of valid blocks for that input.","commonSituations":"Multiple inputs with different sizes but a grid sized for the largest; non-divisible shapes without padding; index_map that scales program_id beyond the valid block count for one specific input.","solutions":["Ensure every input's BlockSpec yields exactly ceil(input.shape / block_shape) valid blocks matching the grid","Pad smaller inputs so their block counts align with the grid","Fix the index_map to clamp or correctly compute per-input block offsets","Reproduce in interpret mode and inspect the reported block_indices and read_range vs the input shape"],"exampleFix":"# before\nin_spec = BlockSpec((BM,), lambda i: (i * BM,))  # grid too big for this input\n# after\nin_spec = BlockSpec((BM,), lambda i: (i * BM,))\ngrid = (x.shape[0] // BM,)  # match grid to smallest valid block count (pad x if needed)","handlingStrategy":"validation","validationCode":"import math\ndef check_blockspec(x, spec, grid):\n    n_blocks = math.ceil(x.shape[0] / spec.block_shape[0])\n    assert grid[0] <= n_blocks, f'grid {grid} exceeds {n_blocks} blocks for input shape {x.shape}'","typeGuard":null,"tryCatchPattern":"try:\n    kernel_grid_run(inputs)\nexcept IndexError as e:\n    if 'Out-of-bounds block index' in str(e) and 'input' in str(e):\n        pad_and_retry()","preventionTips":["Derive grid from the smallest valid block count across all inputs","Pad each input so its block count matches the grid","Log BlockSpec index_map outputs during development"],"tags":["jax","pallas","tpu","blockspec","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"}