jax-ml/jax · error · IndexError

index {i} out of bounds for axis {axis} with size {size} ({n

Error message

index {i} out of bounds for axis {axis} with size {size} ({normalize_indices=})

What it means

validate_static_indices checks static integer indices against the array's static shape; after optional negative-index normalization, the index must satisfy 0 <= i < size for that axis. This check runs when mode='promise_in_bounds' paths (e.g. to_static_slice) validate user indices.

Source

Thrown at jax/_src/numpy/indexing.py:214

    """Create an NDIndexer object from raw user-supplied indices."""
    indices = eliminate_deprecated_list_indexing(indices)
    parsed = _parse_indices(indices, shape)
    return cls(shape=shape, indices=parsed)

  def validate_static_indices(self, normalize_indices: bool = True) -> None:
    """Check that all static integer indices are in-bounds.

    Raises an IndexError in case of out-of-bound indices
    """
    for idx in self.indices:
      if idx.typ == IndexType.INTEGER:
        assert isinstance(idx.index, (int, np.integer))
        i = operator.index(idx.index)
        axis, = idx.consumed_axes
        size = self.shape[axis]
        normed_idx = i + size if normalize_indices and i < 0 else i
        if not 0 <= normed_idx < size:
          raise IndexError(f"index {i} out of bounds for axis {axis} with size {size}"
                           f" ({normalize_indices=})")

  def validate_slices(self) -> None:
    """Check that all slices have static start/stop/step values.

    Raises an IndexError in case of non-static entries.
    """
    for position, idx in enumerate(self.indices):
      if idx.typ == IndexType.SLICE:
        assert isinstance(idx.index, slice)
        elts = [idx.index.start, idx.index.stop, idx.index.step]
        if not all(_is_slice_element_none_or_constant_or_symbolic(val)
                   for val in elts):
          msg = ("Array slice indices must have static start/stop/step to be used "
                 f"with NumPy indexing syntax. Got {idx.index} at position "
                 f"{position}. To index an array at a dynamic position with a "
                 "static slice size, use x[jax.ds(start, size)] or "
                 "lax.dynamic_slice/dynamic_update_slice instead (JAX does not "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp or modulo the index: x[i % x.shape[0]]
2. Re-check the shape at runtime with x.shape[axis] before indexing
3. Fix off-by-one in loop ranges (range(n) vs range(n+1))

Example fix

// before
y = x[i]  # i may equal n on last step
// after
y = x[min(i, x.shape[0] - 1)]

n = x.shape[axis]
i = i + n if i < 0 else i
assert 0 <= i < n, f'{i} not in [0,{n})'

try:
    y = x[i]
except IndexError:
    y = x[i % x.shape[0]]  # wrap-around fallback

When it happens

Trigger: x = jnp.zeros(3); x[5] or x[-4]; also x[3] on an axis whose size shrank after a reshape/config change. Raised in static-index validation during rewriting_take/to_static_slice.

Common situations: Off-by-one loop bounds, hardcoded index for a shape that changed (batch-size config), negative index equal to -size-1, or empty axis edge cases.

Related errors

• index is out of bounds for axis {axis} with size 0
• Value of type {type(self)} is not convertible to integer ind
• TracerIntegerConversionError
• reduction axes {axes} contains out-of-bounds indices for {op
• axis argument out of range: {axis=} for {operand.shape=}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/cf9541eb4d14d8d0. Report an issue: GitHub.