{"record":{"id":"15570a453ad63da0","repo":"jax-ml/jax","slug":"slice-idx-along-axis-axis-is-out-of-bounds-for","errorCode":null,"errorMessage":"Slice {idx} along axis {axis} is out of bounds for shape {shape}","messagePattern":"Slice (.+?) along axis (.+?) is out of bounds for shape (.+?)","errorType":"validation","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":976,"sourceCode":"        raise IndexError(\n            f\"Index {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(idx if idx >= 0 else bound + idx)\n      slice_shape.append(1)\n      is_squeezed.append(True)\n    elif isinstance(idx, slice):\n      if idx.step is not None and idx.step != 1:\n        raise NotImplementedError(\"Strided slices not implemented\")\n      start = idx.start or 0\n      if start < 0:\n        start = bound + start\n      stop = idx.stop or bound\n      if stop < 0:\n        stop = bound + stop\n      if check_oob and (\n          start < 0 or start >= bound or stop < 0 or stop > bound\n      ):\n        raise IndexError(\n            f\"Slice {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(start)\n      slice_shape.append(stop - start)\n      is_squeezed.append(False)\n    elif isinstance(idx, DynamicSlice):\n      if check_oob and (\n          isinstance(idx.base, int) and idx.base + idx.length > bound\n      ):\n        raise IndexError(\n            f\"Slice {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(idx.base)\n      slice_shape.append(idx.length)\n      is_squeezed.append(False)\n    elif isinstance(idx, ir.Value):\n      if not isinstance(idx.type, ir.IndexType):\n        raise ValueError(\"Expected an index-typed index\")","sourceCodeStart":958,"sourceCodeEnd":994,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L958-L994","documentation":"Slice bounds on a memref axis are validated: after normalizing negative start/stop, the slice must satisfy 0 <= start < bound and 0 <= stop <= bound. Otherwise Mosaic raises IndexError with the offending slice, like numpy would.","triggerScenarios":"ref[10:20] on an axis of size 8; ref[5:-10] where -10 normalizes below 0; ref[start:] with start computed from a loop variable exceeding the bound; with check_oob enabled.","commonSituations":"Edge tiles in a tiled kernel where the tile slice extends past the tensor extent; reusing slice parameters after changing the buffer shape; forgetting that unlike numpy, Mosaic does not silently clamp stop to bound.","solutions":["Clamp the slice: stop = min(stop, bound); start = max(0, min(start, bound)) before indexing","Pad the source tensor so slice extents divide evenly into tiles","Compute slices from the actual shape: ir.MemRefType(ref.type).shape[axis]"],"exampleFix":"# before\nedge = buf[off : off + BLOCK]  # off+BLOCK > size on last tile\n# after\nsize = ir.MemRefType(buf.type).shape[0]\nedge = buf[off : min(off + BLOCK, size)]","handlingStrategy":"validation","validationCode":"bound = ir.MemRefType(ref.type).shape[axis]\nstart = max(0, s.start or 0)\nstop = max(start, min(s.stop if s.stop is not None else bound, bound))\nassert 0 <= start < bound and start <= stop <= bound","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Clamp slice stop to the axis extent for edge tiles","Compute all slices from ir.MemRefType(...).shape, not assumed sizes"],"tags":["jax","mosaic-gpu","memref","slice","index-out-of-bounds"],"backgroundTag":"slice-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}