{"record":{"id":"fcf3acb2fed26445","repo":"jax-ml/jax","slug":"index-idx-along-axis-axis-is-out-of-bounds-for","errorCode":null,"errorMessage":"Index {idx} along axis {axis} is out of bounds for shape {shape}","messagePattern":"Index (.+?) along axis (.+?) is out of bounds for shape (.+?)","errorType":"validation","errorClass":"IndexError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":958,"sourceCode":"  )\n\n\ndef parse_indices(\n    index, shape: Sequence[int], *, check_oob: bool = True\n) -> tuple[list[ir.Value | int], list[int], list[bool]]:\n  if not isinstance(index, tuple):\n    index = (index,)\n  if trailing_dims := len(shape) - len(index):\n    index += (slice(None),) * trailing_dims\n  base_indices: list[ir.Value | int] = []\n  slice_shape = []\n  is_squeezed = []\n  for axis, (idx, bound) in enumerate(zip(index, shape)):\n    if isinstance(idx, (ir.Operation, ir.OpView)):\n      idx = idx.result\n    if isinstance(idx, int):\n      if check_oob and (idx >= bound or (idx < 0 and -idx > bound)):\n        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(","sourceCodeStart":940,"sourceCodeEnd":976,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L940-L976","documentation":"When indexing a Mosaic memref via __getitem__/__setitem__ or memref_slice, integer indices are bounds-checked against the axis extent. An index >= bound (or negative beyond -bound) raises this IndexError, mirroring numpy semantics for GPU memref access in kernels.","triggerScenarios":"ref[idx] where idx is a Python int >= shape[axis] or < -shape[axis] with check_oob enabled; e.g. ref[64] on a shape (32, 8) memref, or ref[-33] where bound is 32.","commonSituations":"Using grid/block indices from launch geometry to index a shared-memory tile without clamping; shrinking a test tensor without updating indices; converting numpy prototype code that used out-of-bounds-but-never-hit indices.","solutions":["Clamp or guard indices with min/max against the axis extent before indexing","Verify the memref shape with ir.MemRefType(ref.type).shape in debug builds and fix hardcoded indices","Pass check_oob=False only if you can prove (e.g. predicated execution) the index is never dereferenced"],"exampleFix":"# before\nx = tile[row, col]  # col can be >= tile width on edge blocks\n# after\nfrom jax.experimental.mosaic.gpu import utils\ncol_c = min(col, width - 1) if isinstance(col, int) else col\nx = tile[row, col_c]","handlingStrategy":"validation","validationCode":"shape = ir.MemRefType(ref.type).shape\nassert all(-shape[a] <= i < shape[a] for a, i in enumerate(idx_tuple) if isinstance(i, int))","typeGuard":"def in_bounds(shape, idx) -> bool:\n    return all(-b <= i < b for b, i in zip(shape, idx) if isinstance(i, int))","tryCatchPattern":"try:\n    v = tile[i, j]\nexcept IndexError:\n    v = zero  # predicated-off lane\n    # or: continue","preventionTips":["Clamp edge-tile indices against the axis extent","Keep tensor shapes and hardcoded indices in one config constant"],"tags":["jax","mosaic-gpu","memref","index-out-of-bounds","indexerror"],"backgroundTag":"index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}