{"record":{"id":"c35f50b3711edf39","repo":"jax-ml/jax","slug":"out-of-bounds-block-index-block-indices-for-outp","errorCode":null,"errorMessage":"Out-of-bounds block index {block_indices} for output \"{output_name}\" in iteration {grid_loop_idx} on device {device_id} (core {local_core_id}): reading [{write_range}] but output has shape {shape}.","messagePattern":"Out-of-bounds block index (.+?) for output \"(.+?)\" in iteration (.+?) on device (.+?) \\(core (.+?)\\): reading \\[(.+?)\\] but output has shape (.+?)\\.","errorType":"validation","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":788,"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  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:\n      src_local_core_id = local_core_id\n    assert races is not None\n    assert clock is not None\n    races.check_write(\n        (src_device_id, src_local_core_id),\n        clock,\n        (memory_space, buffer_id, device_id, local_core_id_for_buffer),\n        write_range,","sourceCodeStart":770,"sourceCodeEnd":806,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L770-L806","documentation":"While copying a kernel buffer's result back into a block of a pallas_call output after kernel invocation, the interpreter found the block index for that output iteration reads outside the output's shape. It names the output, grid iteration, device, and core where the violation occurred.","triggerScenarios":"An output BlockSpec whose block shape/index_map addresses beyond out.shape on some grid iteration — typically a grid larger than the number of output blocks, or a non-divisible output size without padding.","commonSituations":"Different input/output sizes with one shared grid; output tensors smaller than the block shape; index_map bugs scaling program_id too far for the output.","solutions":["Align the grid with the output's block count (ceil(out.shape / block_shape))","Pad outputs to multiples of the block shape","Fix the output index_map so start + block_shape <= out.shape on every iteration"],"exampleFix":"# before\nout_spec = BlockSpec((BN,), lambda i: (i * BN,))\ngrid = (N_total // BM,)  # sized for input, too big for output\n# after\ngrid = (min(x.shape[0], out.shape[0]) // BM,)  # or pad both to a common multiple","handlingStrategy":"validation","validationCode":"import math\nn_out_blocks = math.ceil(out.shape[0] / out_spec.block_shape[0])\nassert grid[0] <= n_out_blocks, f'grid {grid} > output blocks {n_out_blocks}'","typeGuard":null,"tryCatchPattern":"try:\n    run(inputs)\nexcept IndexError as e:\n    if 'Out-of-bounds block index' in str(e) and 'output' in str(e):\n        pad_output_and_retry()","preventionTips":["Derive grid from output as well as input block counts","Keep input/output shapes consistent or padded to common block multiples","Verify output index_map in interpret mode before hardware runs"],"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"}