jax-ml/jax · error · TypeError

dynamic_slice: unrecognized index {pidx.index} at position {

Error message

dynamic_slice: unrecognized index {pidx.index} at position {position}.

What it means

A catch-all TypeError for an index whose classified IndexType does not match any known category (integer, slice, ellipsis, none, array, boolean, dynamic_slice) during dynamic-slice lowering.

Source

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

      elif pidx.typ == IndexType.DYNAMIC_SLICE:
        assert isinstance(pidx.index, indexing.Slice)
        if pidx.index.stride != 1:
          raise TypeError("dynamic_slice: only unit steps supported in slice."
                          f" Got {pidx.index} at position {position}")
      elif pidx.typ == IndexType.SLICE:
        assert isinstance(pidx.index, slice)
        if pidx.index.step is not None and pidx.index.step not in [-1, 1]:
          raise TypeError("dynamic_slice: only unit steps supported in slice."
                          f" Got {pidx.index} at position {position}")
      elif pidx.typ == IndexType.ARRAY:
        if isinstance(pidx.index, Sequence) or np.shape(pidx.index) != ():  # pyrefly: ignore[no-matching-overload]
          raise TypeError("dynamic_slice: only scalar indices allowed."
                          f" Got index of type {type(pidx.index)} at position {position}")
      elif pidx.typ == IndexType.BOOLEAN:
        raise TypeError("dynamic_slice: indices must be scalars or slices."
                        f" Got index of type {type(pidx.index)} at position {position}")
      else:
        raise TypeError(f"dynamic_slice: unrecognized index {pidx.index} at position {position}.")

    start_indices: list[ArrayLike] = []
    slice_sizes: list[int] = []
    rev_axes: list[int] = []
    squeeze_axes: list[int] = []
    newaxis_dims: list[int] = []

    expanded = self.expand_ellipses()
    trivial_slicing = True
    for pidx in expanded.indices:
      if pidx.typ in [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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Stick to plain int, slice, None, Ellipsis, and scalar tracer indices
  2. Convert custom index objects to int or jnp scalar before indexing
  3. Report/check JAX version if using jax.indexing constructs that should be supported
Defensive patterns

Strategy: validation

Validate before calling

allowed = (int, slice, type(None), type(Ellipsis))
assert all(i is Ellipsis or i is None or isinstance(i, allowed) or hasattr(i, 'dtype') or hasattr(i, '__index__') for i in idx_tuple)

Prevention

When it happens

Trigger: Passing an exotic/unsupported index object (custom class with __index__, partially-traced objects, or an internal IndexType added without dynamic-slice support) to dynamic-slice-based indexing.

Common situations: Custom index-like objects passed to x.at[...]; version mismatches where new index types aren't supported by the dynamic path; rarely hit in normal use.

Related errors

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