jax-ml/jax · error · IndexError

Unrecognized index type: {typ}

Error message

Unrecognized index type: {typ}

What it means

_parse_indices hit an IndexType enum value it does not know when computing how many dimensions each index consumes. In practice this is an internal invariant violation, not something user input should produce — the from_index classifier only emits the known types.

Source

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

  ellipses_indices: list[int] = []
  index_types: list[IndexType] = []
  for i, idx in enumerate(indices):
    typ = IndexType.from_index(idx)
    index_types.append(typ)

    if typ == IndexType.NONE:
      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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure a single consistent jax version is installed (pip install -U jax)
  2. Don't construct ParsedIndex/IndexType by hand; go through public x[...] API or from_raw_indices with raw indices
  3. If it appears with normal indexing, report a JAX bug with a minimal repro
Defensive patterns

Strategy: validation

Validate before calling

import jax; assert jax.__version__ == jax.lib.version  # version consistency check

Prevention

When it happens

Trigger: Calling NDIndexer.from_raw_indices / internal parsing with a manually constructed index list containing an out-of-enum IndexType, or a mismatched JAX version where new enum members aren't handled.

Common situations: Mixing JAX versions in one environment, monkey-patching internals, or constructing IndexType/ParsedIndex objects directly in downstream libraries.

Related errors


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