jax-ml/jax · error · IndexError

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

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.

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)

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
Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

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


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