jax-ml/jax · error · ValueError

static_slice with partial slices does not support nontrivial

Error message

static_slice with partial slices does not support nontrivial array sharding.

What it means

When an array is sharded across multiple devices (e.g. named sharding / sharded jitted arrays) and the index expression contains partial slices (like x[:2] instead of full x[:]), lowering to a static slice would produce incorrect sharding semantics, so to_static_slice raises this ValueError.

Source

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

    else:
      parsed_mode = slicing.GatherScatterMode.from_any(mode)
    if any(core.is_symbolic_dim(s) for s in self.shape):
      raise ValueError("mode='slice' is not valid for polymorphic shapes.")

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

    # Validation of the unmodified user indices.
    if parsed_mode == slicing.GatherScatterMode.PROMISE_IN_BOUNDS:
      self.validate_static_indices(normalize_indices=normalize_indices)
    self.validate_slices()

    # 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("static_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.SLICE, IndexType.NONE]:
        pass
      elif pidx.typ in [IndexType.ARRAY, IndexType.BOOLEAN, IndexType.DYNAMIC_SLICE]:
        raise TypeError("static_slice: indices must be static scalars or slices."
                        f" Got index of type {type(pidx.index)} at position {position}")
      else:
        raise TypeError(f"static_slice: unrecognized index {pidx.index} at position {position}.")

    # Now re-iterate to generate static slices.
    start_indices: list[int] = []
    limit_indices: list[int] = []
    strides: list[int] = []
    rev_axes: list[int] = []
    squeeze_axes: list[int] = []
    newaxis_dims: list[int] = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the slice full along sharded axes: x[:, :2] if axis 0 is sharded, or use x[:] then follow-up ops
2. Convert to a single-device array first: jax.device_put(x, jax.devices('cpu')[0]) before slicing
3. Shard a different axis or keep the sliced axis unsharded in the PartitionSpec

Example fix

// before
sh = jax.sharding.NamedSharding(mesh, P('data', None))
x = jax.device_put(big, sh)
y = x[:2]
// after
y = jax.device_put(x, jax.devices('cpu')[0])[:2]

def sharded(x):
    return isinstance(x, jax.Array) and x.sharding.num_devices > 1
if sharded(x):
    x = jax.device_put(x, jax.devices('cpu')[0])

When it happens

Trigger: Indexing a multi-device sharded jax.Array with a partial slice: sharded_arr[:2], sharded_arr[0], inside or outside jitted code using sharding constraints.

Common situations: Moving to jax.sharding / multi-GPU or TPU sharding and reusing single-device slicing code; SPMD code with PartitionSpec where an axis is split and then sliced.

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/a9d30f2d16acd43d. Report an issue: GitHub.