{"record":{"id":"371f8952bf711666","repo":"jax-ml/jax","slug":"dynamic-slice-with-partial-slices-does-not-support","errorCode":null,"errorMessage":"dynamic_slice with partial slices does not support nontrivial array sharding.","messagePattern":"dynamic_slice with partial slices does not support nontrivial array sharding\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":487,"sourceCode":"      self, *,\n      arr_is_sharded: bool = False,\n      normalize_indices: bool = True,\n      mode: str | slicing.GatherScatterMode | None) -> _DynamicSliceIndexer:\n    \"\"\"Convert to DynamicSliceIndexer data structure.\n\n    If this is not possible, raise a ValueError, TypeError, or IndexError.\n    \"\"\"\n    if mode is not None:\n      parsed_mode = slicing.GatherScatterMode.from_any(mode)\n      if parsed_mode not in [\n          slicing.GatherScatterMode.PROMISE_IN_BOUNDS, slicing.GatherScatterMode.CLIP]:\n        raise ValueError(\"dynamic_slice requires mode='promise_in_bounds' or mode='clip'\")\n\n    # For sharded inputs, indexing (like x[0]) and partial slices (like x[:2] as\n    # opposed to x[:]) lead to incorrect sharding semantics when computed via slice.\n    # TODO(yashkatariya): fix slice with sharding\n    if arr_is_sharded and self.has_partial_slices():\n      raise ValueError(\"dynamic_slice with partial slices does not support nontrivial array sharding.\")\n\n    for position, pidx in enumerate(self.indices):\n      if pidx.typ in [IndexType.INTEGER, IndexType.ELLIPSIS, IndexType.NONE]:\n        pass\n      elif pidx.typ == IndexType.DYNAMIC_SLICE:\n        assert isinstance(pidx.index, indexing.Slice)\n        if pidx.index.stride != 1:\n          raise TypeError(\"dynamic_slice: only unit steps supported in slice.\"\n                          f\" Got {pidx.index} at position {position}\")\n      elif pidx.typ == IndexType.SLICE:\n        assert isinstance(pidx.index, slice)\n        if pidx.index.step is not None and pidx.index.step not in [-1, 1]:\n          raise TypeError(\"dynamic_slice: only unit steps supported in slice.\"\n                          f\" Got {pidx.index} at position {position}\")\n      elif pidx.typ == IndexType.ARRAY:\n        if isinstance(pidx.index, Sequence) or np.shape(pidx.index) != ():  # pyrefly: ignore[no-matching-overload]\n          raise TypeError(\"dynamic_slice: only scalar indices allowed.\"\n                          f\" Got index of type {type(pidx.index)} at position {position}\")","sourceCodeStart":469,"sourceCodeEnd":505,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L469-L505","documentation":"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.","triggerScenarios":"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)].","commonSituations":"SPMD/multi-GPU refactors of single-device code; jax.jit with in_shardings splitting an axis that is then sliced or element-picked.","solutions":["Avoid indexing/slicing the sharded axis; operate on full axes","Reshard to one device before slicing: jax.device_put(x, single_device_sharding)[i]","Adjust the PartitionSpec so the sliced axis is replicated"],"exampleFix":"// before\nx = jax.device_put(big, NamedSharding(mesh, P('dev', None)))\ny = x[0]\n// after\ny = jax.device_put(x, jax.sharding.SingleDeviceSharding(jax.devices()[0]))[0]","handlingStrategy":"fallback","validationCode":"if getattr(x, 'sharding', None) and x.sharding.num_devices > 1:\n    x = jax.device_put(x, jax.sharding.SingleDeviceSharding(jax.devices()[0]))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't partial-slice sharded axes; move data or reshard first","Design PartitionSpecs so sliced axes stay replicated"],"tags":["jax","sharding","dynamic-slice","multi-device"],"backgroundTag":"sharded-array-slicing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}