jax-ml/jax · error · TypeError

dynamic_slice: only scalar indices allowed. Got index of typ

Error message

dynamic_slice: only scalar indices allowed. Got index of type {type(pidx.index)} at position {position}

What it means

In dynamic-slice lowering, indices classified as ARRAY type must be scalar (0-dimensional). If a sequence or non-scalar array is used as an index, JAX raises this TypeError because dynamic_slice only takes scalar start indices.

Source

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

    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] = []
    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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use standard fancy indexing (x[np.array([i, j])]) instead of the dynamic-slice path
Squeeze the index to a scalar: index = jnp.asarray(index).reshape(()) when it holds one element
Use lax.gather or x.at[indices_array].get() for multiple indices

Example fix

// before
idx = [i]
y = x.at[idx].get()  # sequence index
// after
y = x.at[i].get()  # scalar index

Defensive patterns

Strategy: type-guard

Validate before calling

import numpy as np
assert all(np.shape(i) == () and not isinstance(i, (list, tuple)) for i in idx_tuple if not isinstance(i, slice))

Type guard

def is_scalar_index(i) -> bool:
    import numpy as np
    return not isinstance(i, (list, tuple)) and np.shape(i) == ()

Try / catch

try:
    y = x.at[idx].get()
except TypeError as e:
    if 'only scalar indices allowed' in str(e):
        y = x[np.asarray(idx)]  # fall back to gather-style indexing
    else:
        raise

Prevention

Use scalar tracers or ints in dynamic-slice positions
Use array indexing (x[np.array([...])]) for multiple indices

When it happens

Trigger: Passing a list, tuple, or shape-(n,) array as an index in a dynamic index expression, e.g. x.at[[0, 1]].get() routed through to_dynamic_slice, or mixing gather-style arrays into dynamic slicing.

Common situations: Mixing NumPy fancy-indexing habits (arr[[i, j]]) with JAX dynamic slicing; passing an index array where a scalar tracer was expected in jit code; index variables accidentally wrapped in lists.

Related errors

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