jax-ml/jax · error · TypeError

dynamic_slice: indices must be scalars or slices. Got index

Error message

dynamic_slice: indices must be scalars or slices. Got index of type {type(pidx.index)} at position {position}

What it means

Dynamic slicing only accepts integers and slices. Boolean indices (and any index typed BOOLEAN) are rejected here because dynamic_slice cannot represent boolean mask gathers.

Source

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

    for position, pidx in enumerate(self.indices):
      if pidx.typ in [IndexType.INTEGER, IndexType.ELLIPSIS, IndexType.NONE]:
        pass
      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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert booleans to integers: index = jnp.asarray(mask).astype(jnp.int32) if it encodes positions
Use proper boolean masking x[mask] outside the dynamic-slice path
Use jnp.nonzero(mask) to convert a mask to integer indices first

Example fix

// before
y = x.at[some_bool].get()
// after
y = x.at[int(some_bool)].get()  # or use x[x > 0] for masks

Defensive patterns

Strategy: type-guard

Validate before calling

assert not any(isinstance(i, (bool, np.bool_)) or (hasattr(i, 'dtype') and i.dtype == np.bool_) for i in idx_tuple), 'boolean indices unsupported in dynamic slice'

Type guard

def is_bool_index(i) -> bool:
    return isinstance(i, (bool, np.bool_)) or (hasattr(i, 'dtype') and getattr(i.dtype, 'name', '') == 'bool')

Prevention

Convert masks with jnp.nonzero(mask) before indexing
Never pass comparison results as positional indices

When it happens

Trigger: Using a boolean array or scalar as an index in the dynamic-slice path, e.g. x.at[bool_tracer].get() or passing a boolean mask where an integer index is expected.

Common situations: Reusing a comparison result (arr > 0) as an index in code migrated to dynamic indexing; bool flags accidentally used as positional indices; masks intended for x[mask] boolean indexing.

Related errors

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