jax-ml/jax · error · ValueError

dynamic_slice with partial slices does not support nontrivia

Error message

dynamic_slice with partial slices does not support nontrivial array sharding.

What it means

Dynamic counterpart of error 1433: when the array is sharded across devices and the index expression contains partial slices, lowering via dynamic slicing would give wrong sharding semantics, so to_dynamic_slice raises ValueError.

Source

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

      self, *,
      arr_is_sharded: bool = False,
      normalize_indices: bool = True,
      mode: str | slicing.GatherScatterMode | None) -> _DynamicSliceIndexer:
    """Convert to DynamicSliceIndexer data structure.

    If this is not possible, raise a ValueError, TypeError, or IndexError.
    """
    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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid indexing/slicing the sharded axis; operate on full axes
2. Reshard to one device before slicing: jax.device_put(x, single_device_sharding)[i]
3. Adjust the PartitionSpec so the sliced axis is replicated

Example fix

// before
x = jax.device_put(big, NamedSharding(mesh, P('dev', None)))
y = x[0]
// after
y = jax.device_put(x, jax.sharding.SingleDeviceSharding(jax.devices()[0]))[0]

if getattr(x, 'sharding', None) and x.sharding.num_devices > 1:
    x = jax.device_put(x, jax.sharding.SingleDeviceSharding(jax.devices()[0]))

When it happens

Trigger: Indexing a sharded jax.Array (multi-device, PartitionSpec splitting an axis) with a partial slice or dynamic slice on the sharded axis, e.g. sharded_x[0] or sharded_x[jax.ds(i, k)].

Common situations: SPMD/multi-GPU refactors of single-device code; jax.jit with in_shardings splitting an axis that is then sliced or element-picked.

Related errors

• to_dlpack can only pack a dlpack tensor from an array on a s
• Query, key and value should have same sharding.
• from_dlpack can only unpack a dlpack tensor onto a singular
• 0th dimension of all xs should be replicated. Got {}
• 0th dimension of leaf passed to `jax.lax.map` should be repl

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