jax-ml/jax · error · IndexError

index is out of bounds for axis

Error message

index is out of bounds for axis {axis} with size 0

What it means

to_static_slice explicitly rejects integer indexing into an axis whose size is statically 0, because the underlying XLA slice op would error or produce invalid results on empty axes. Even though the index is a scalar, selecting from an empty axis is out of bounds.

Solutions

1. Guard on shape before indexing: if x.shape[0]: ...
2. Use x[:1] (slice) which safely yields an empty result on empty axes
3. Fix upstream filtering so the empty case is handled explicitly

Example fix

// before
first = x[0]
// after
first = x[:1]  # empty-safe

if x.shape[axis] == 0:
    raise ValueError(f'axis {axis} is empty')

When it happens

Trigger: x = jnp.zeros((0, 4)); x[0] — any integer index on a statically-empty axis via the static-slice path.

Common situations: Pipeline stages where a batch/filter step legitimately produces zero rows and downstream code unconditionally indexes row 0; empty splits from data filtering.

Related errors

• Cannot do a non-empty jnp.take() from an empty axis.
• index out of bounds for axis with size ( )
• an index can only have a single ellipsis ('...')
• argmin and argmax require non-empty reduced dimension…
• attempt to get argmax of an empty sequence

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

Appendix: source

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

    squeeze_axes: list[int] = []
    newaxis_dims: list[int] = []

    expanded = self.expand_ellipses()
    for pidx in expanded.indices:
      if pidx.typ in [IndexType.ARRAY, IndexType.BOOLEAN, IndexType.ELLIPSIS]:
        raise RuntimeError(f"Internal: unexpected index encountered: {pidx}")
      elif pidx.typ == IndexType.NONE:
        # Expanded axes indices are based on the rank of the array after slicing
        # (tracked by start_indices) and squeezing (tracked by squeeze_axes), and
        # expand_dims inserts dimensions in order, so we must also account for
        # previous expanded dimensions.
        newaxis_dims.append(len(start_indices) - len(squeeze_axes) + len(newaxis_dims) )
      elif pidx.typ == IndexType.INTEGER:
        assert isinstance(pidx.index, (int, np.integer))
        axis, = pidx.consumed_axes
        if core.definitely_equal(self.shape[axis], 0):
          # XLA gives error when indexing into an axis of size 0
          raise IndexError(f"index is out of bounds for axis {axis} with size 0")
        start_index = int(pidx.index)
        if normalize_indices and start_index < 0:
          start_index += self.shape[axis]
        # Normalization & validation have already been handled, so clip start_index
        # to valid range
        start_index = min(max(start_index, 0), self.shape[axis] - 1)
        start_indices.append(start_index)
        limit_indices.append(start_index + 1)
        strides.append(1)
        squeeze_axes.append(axis)
      elif pidx.typ == IndexType.SLICE:
        assert isinstance(pidx.index, slice)
        axis, = pidx.consumed_axes
        size = self.shape[axis]
        start, stop, stride = pidx.index.indices(size)
        if stride < 0:
          new_start = min(size, stop + 1 + abs(start - stop - 1) % abs(stride))
          start_indices.append(new_start)

View on GitHub (pinned to 1e1c6a8fc0)