jax-ml/jax · error · IndexError

Array slice indices must have static start/stop/step to be u

Error message

Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Got {idx.index} at position {position}. To index an array at a dynamic position with a static slice size, use x[jax.ds(start, size)] or lax.dynamic_slice/dynamic_update_slice instead (JAX does not support dynamically sized arrays within traced functions).

What it means

validate_slices rejects slice objects whose start/stop/step are non-static (contain Tracers) when NumPy slice syntax is used under tracing (jit/vmap/grad). JAX needs static slice bounds because array shapes must be static inside traces; the error points to jax.ds(start, size) / lax.dynamic_slice as the alternative.

Source

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

    Raises an IndexError in case of non-static entries.
    """
    for position, idx in enumerate(self.indices):
      if idx.typ == IndexType.SLICE:
        assert isinstance(idx.index, slice)
        elts = [idx.index.start, idx.index.stop, idx.index.step]
        if not all(_is_slice_element_none_or_constant_or_symbolic(val)
                   for val in elts):
          msg = ("Array slice indices must have static start/stop/step to be used "
                 f"with NumPy indexing syntax. Got {idx.index} at position "
                 f"{position}. To index an array at a dynamic position with a "
                 "static slice size, use x[jax.ds(start, size)] or "
                 "lax.dynamic_slice/dynamic_update_slice instead (JAX does not "
                 "support dynamically sized arrays within traced functions).")
          tracer = next((val for val in elts if isinstance(val, core.Tracer)), None)
          if tracer is not None:
            msg += tracer._origin_msg()
          raise IndexError(msg)

  @staticmethod
  def is_sharded(arr) -> bool:
    """Check whether the array is sharded."""
    return isinstance(arr, array.ArrayImpl) and not arr.sharding.num_devices == 1

  def has_partial_slices(self) -> bool:
    """Check whether the indexer contains partial slices.

    For sharded arrays, partial slices cannot automatically propagate
    sharding.
    """
    for idx in self.indices:
      if idx.typ in [IndexType.INTEGER, IndexType.DYNAMIC_SLICE]:
        return True
      if idx.typ == IndexType.SLICE:
        slc = idx.index
        assert isinstance(slc, slice)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.lax.dynamic_slice(x, (start,), (size,)) with a fixed size
2. Use the jax.ds(start, size) dynamic-slice index marker in x[...]
3. Keep slice bounds static: hoist i out of jit or close over Python ints

Example fix

# before
@jax.jit
def f(x, i):
    return x[i:i+5]
# after
@jax.jit
def f(x, i):
    return jax.lax.dynamic_slice(x, (i,), (5,))

When it happens

Trigger: @jit-def f(x): return x[i:i+5] where i is a traced value (argument or computed from one); any slice built from Tracer values inside jit/vmap/scan.

Common situations: Sliding windows at runtime offsets, batching with dynamic start positions, porting NumPy windowing code under jit. Very common when enabling jit on existing code.

Related errors

• The unsafe_buffer_pointer() method was called on {self._erro
• function {dbg.func_src_info} traced for {dbg.traced_for} ret
• When used within transformed code, jax.experimental.random.s
• Formatting arguments to checkify.check need to be PyTrees of
• Value of type {type(self)} is not convertible to float.

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