jax-ml/jax · error · IndexError

Too many indices: array is {len(shape)}-dimensional, but {to

Error message

Too many indices: array is {len(shape)}-dimensional, but {total_consumed} were indexed

What it means

The index expression consumes more dimensions than the array has: each integer/slice/array index consumes one axis and each boolean array consumes its ndim axes; the total exceeded len(shape). Matches NumPy's 'too many indices for array' error.

Source

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

      dimensions_consumed.append(0)
    elif typ == IndexType.ELLIPSIS:
      # We don't yet know how many dimensions are consumed, so set to zero
      # for now and update later.
      dimensions_consumed.append(0)
      ellipses_indices.append(i)
    elif typ == IndexType.BOOLEAN:
      dimensions_consumed.append(np.ndim(idx))  # pyrefly: ignore[bad-argument-type]
    elif typ in [IndexType.INTEGER, IndexType.ARRAY, IndexType.SLICE, IndexType.DYNAMIC_SLICE]:
      dimensions_consumed.append(1)
    else:
      raise IndexError(f"Unrecognized index type: {typ}")

  # 2. Validate the consumed dimensions and ellipses.
  if len(ellipses_indices) > 1:
    raise IndexError("an index can only have a single ellipsis ('...')")
  total_consumed = sum(dimensions_consumed)
  if total_consumed > len(shape):
    raise IndexError(f"Too many indices: array is {len(shape)}-dimensional,"
                     f" but {total_consumed} were indexed")
  if ellipses_indices:
    dimensions_consumed[ellipses_indices[0]] = len(shape) - total_consumed

  # 3. Generate the final sequence of parsed indices.
  result: list[ParsedIndex] = []
  current_dim = 0
  for index, typ, n_consumed in safe_zip(indices, index_types, dimensions_consumed):
    consumed_axes = tuple(range(current_dim, current_dim + n_consumed))
    current_dim += len(consumed_axes)
    result.append(ParsedIndex(index=index, typ=typ, consumed_axes=consumed_axes))
  return result


@register_pytree_node_class
@dataclasses.dataclass(frozen=True, kw_only=True, slots=True)
class NDIndexer:
  """Object that implements NumPy-style indexing operations on top of JAX.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check x.ndim and the length of your index tuple before indexing
2. Use an ellipsis to absorb extra axes: x[0, ...]
3. Audit reshapes/squeezes upstream that reduced rank unexpectedly

Example fix

// before
y = x[i, j, k]  # x is 2-D
// after
y = x[i, j]  # or x[i, j, ...] when rank varies

assert len(nonnewaxis_idx) <= x.ndim, (x.ndim, nonnewaxis_idx)
# account for boolean masks consuming multiple axes

When it happens

Trigger: x = jnp.zeros((3,4)); x[0, 0, 0] — three indices on a 2-D array; or a boolean mask of higher rank than the array.

Common situations: Code written for a higher-rank array run on squeezed/reshaped data, loops that append index tuples without tracking rank, or accidental unpacking of coordinates tuples.

Related errors

• boolean index did not match shape of indexed array in index
• {name} was requested to map its argument along axis {axis},
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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