jax-ml/jax · error · TypeError

dynamic_slice: only unit steps supported in slice. Got {pidx

Error message

dynamic_slice: only unit steps supported in slice. Got {pidx.index} at position {position}

What it means

Raised by JAX's dynamic-slice index lowering when a jax.indexing.Slice (dynamic slice index) has a stride other than 1. XLA's dynamic_slice primitive only supports unit steps, so non-unit strides cannot be lowered and JAX rejects them early with a TypeError.

Source

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

    if mode is not None:
      parsed_mode = slicing.GatherScatterMode.from_any(mode)
      if parsed_mode not in [
          slicing.GatherScatterMode.PROMISE_IN_BOUNDS, slicing.GatherScatterMode.CLIP]:
        raise ValueError("dynamic_slice requires mode='promise_in_bounds' or mode='clip'")

    # For sharded inputs, indexing (like x[0]) and partial slices (like x[:2] as
    # opposed to x[:]) lead to incorrect sharding semantics when computed via slice.
    # TODO(yashkatariya): fix slice with sharding
    if arr_is_sharded and self.has_partial_slices():
      raise ValueError("dynamic_slice with partial slices does not support nontrivial array sharding.")

    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] = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use stride 1 (jax.indexing.Slice(start, stop, 1)) and select every other element afterwards if needed
Replace the strided dynamic slice with a static Python slice arr[start:stop:2] when indices are concrete
Use lax.gather or arr.at[...] with an index array computed via arange(start, stop, step)

Example fix

// before
idx = (jax.indexing.Slice(0, 10, 2),)
y = x[idx]
// after
idx = (jax.indexing.Slice(0, 10, 1),)
y = x[idx][:, ::2]

Defensive patterns

Strategy: validation

Validate before calling

from jax._src import indexing
assert isinstance(idx, indexing.Slice) is False or idx.stride == 1, 'stride must be 1'

Prevention

Keep dynamic slice strides at 1; do additional striding after slicing
Use static Python slices when bounds are concrete

When it happens

Trigger: Using jax.lax.dynamic_slice or jax.experimental.array_api / jnp indexing with a jax.indexing.Slice(start, stop, stride) where stride != 1, inside code paths that go through rewriting_take/to_dynamic_slice (e.g. XLA-usable index expressions like jnp.ndarray.at or dynamic indexing APIs).

Common situations: Porting NumPy code that uses arr[start:stop:2] into JAX's dynamic indexing API; using jax.indexing.Slice with a computed step; batched gather code assumed to support strided windows.

Related errors

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