{"record":{"id":"fe79310889bc081b","repo":"jax-ml/jax","slug":"out-of-bounds-block-index-block-indices-for-all","errorCode":null,"errorMessage":"Out-of-bounds block index {block_indices} for {allocation_key}, input \"{input_name}\" in iteration {grid_loop_idx}: reading [{read_range}] but input has shape {shape}.","messagePattern":"Out-of-bounds block index (.+?) for (.+?), input \"(.+?)\" in iteration (.+?): reading \\[(.+?)\\] but input has shape (.+?)\\.","errorType":"exception","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py","lineNumber":480,"sourceCode":") -> np.ndarray:\n  \"\"\"Handles out-of-bounds read based on shared_memory configuration.\"\"\"\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            f\"Out-of-bounds read of {allocation_key}:\"\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 {allocation_key},\"\n            f' input \"{input_name}\" in iteration {grid_loop_idx}:'\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    return uninit_array\n  else:\n    uninit_array[tuple(slice(s) for s in ret.shape)] = ret\n    return uninit_array\n","sourceCodeStart":462,"sourceCodeEnd":498,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py#L462-L498","documentation":"Same out-of-bounds detection as the plain read case, but specialized to input blocks that the interpreter must copy into a buffer before running the kernel body. It reports the block indices, grid loop iteration, requested read range, and the actual input shape.","triggerScenarios":"A kernel in interpret mode indexes an input Block (via a transformed reference) such that for some grid iteration the read range exceeds the input's shape; typical when grid is ceil-div but reads are unmasked.","commonSituations":"Tail iterations of non-evenly-divisible shapes; kernels written assuming padded inputs (e.g. from a pipeline with padded batch dim) run without padding; incorrect Block reference constructed from dmapped user indices.","solutions":["Use masked loads/stores (Block with mask or swap in bounded index computation) for boundary iterations","Pad the input to a multiple of the block size before passing it to the kernel","Verify the Block's dimension maps produce indices within [0, dim) for every grid iteration","Reproduce in interpret mode with tiny shapes to identify the offending iteration printed in the message"],"exampleFix":"# before\nx_block = x[ds[start, bs], :]  # unmasked tail read\n# after\nmask = (start + jnp.arange(bs)) < n\nx_block = x[ds[start, bs], :].mask(mask, 0.0)  # or pad x to multiple of bs","handlingStrategy":"validation","validationCode":"n = x.shape[0]\nassert all((i*bs + bs <= n) or use_mask for i in range(nblocks))","typeGuard":null,"tryCatchPattern":"try:\n    out = kernel(x)\nexcept IndexError as e:\n    if 'Out-of-bounds block index' in str(e):\n        pad or mask inputs","preventionTips":["Pad inputs to multiples of block sizes","Default to masked loads on boundary blocks"],"tags":["pallas","mosaic-gpu","interpret-mode","out-of-bounds","input-block"],"backgroundTag":"buffer-out-of-bounds-read","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}