{"record":{"id":"a9d30f2d16acd43d","repo":"jax-ml/jax","slug":"static-slice-with-partial-slices-does-not-support","errorCode":null,"errorMessage":"static_slice with partial slices does not support nontrivial array sharding.","messagePattern":"static_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":397,"sourceCode":"    else:\n      parsed_mode = slicing.GatherScatterMode.from_any(mode)\n    if any(core.is_symbolic_dim(s) for s in self.shape):\n      raise ValueError(\"mode='slice' is not valid for polymorphic shapes.\")\n\n    if parsed_mode not in [\n        slicing.GatherScatterMode.PROMISE_IN_BOUNDS, slicing.GatherScatterMode.CLIP]:\n      raise ValueError(\"static_slice requires mode='promise_in_bounds' or mode='clip'\")\n\n    # Validation of the unmodified user indices.\n    if parsed_mode == slicing.GatherScatterMode.PROMISE_IN_BOUNDS:\n      self.validate_static_indices(normalize_indices=normalize_indices)\n    self.validate_slices()\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(\"static_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.SLICE, IndexType.NONE]:\n        pass\n      elif pidx.typ in [IndexType.ARRAY, IndexType.BOOLEAN, IndexType.DYNAMIC_SLICE]:\n        raise TypeError(\"static_slice: indices must be static scalars or slices.\"\n                        f\" Got index of type {type(pidx.index)} at position {position}\")\n      else:\n        raise TypeError(f\"static_slice: unrecognized index {pidx.index} at position {position}.\")\n\n    # Now re-iterate to generate static slices.\n    start_indices: list[int] = []\n    limit_indices: list[int] = []\n    strides: list[int] = []\n    rev_axes: list[int] = []\n    squeeze_axes: list[int] = []\n    newaxis_dims: list[int] = []\n","sourceCodeStart":379,"sourceCodeEnd":415,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L379-L415","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make the slice full along sharded axes: x[:, :2] if axis 0 is sharded, or use x[:] then follow-up ops","Convert to a single-device array first: jax.device_put(x, jax.devices('cpu')[0]) before slicing","Shard a different axis or keep the sliced axis unsharded in the PartitionSpec"],"exampleFix":"// before\nsh = jax.sharding.NamedSharding(mesh, P('data', None))\nx = jax.device_put(big, sh)\ny = x[:2]\n// after\ny = jax.device_put(x, jax.devices('cpu')[0])[:2]","handlingStrategy":"fallback","validationCode":"def sharded(x):\n    return isinstance(x, jax.Array) and x.sharding.num_devices > 1\nif sharded(x):\n    x = jax.device_put(x, jax.devices('cpu')[0])","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep sharded axes whole in index expressions","Reshard to single device before ad-hoc slicing/debugging"],"tags":["jax","sharding","multi-device","slicing"],"backgroundTag":"sharded-array-slicing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}