{"record":{"id":"32b68497dce28c91","repo":"jax-ml/jax","slug":"dynamic-slice-only-scalar-indices-allowed-got-in","errorCode":null,"errorMessage":"dynamic_slice: only scalar indices allowed. Got index of type {type(pidx.index)} at position {position}","messagePattern":"dynamic_slice: only scalar indices allowed\\. Got index of type (.+?) at position (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":504,"sourceCode":"    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}\")\n      elif pidx.typ == IndexType.BOOLEAN:\n        raise TypeError(\"dynamic_slice: indices must be scalars or slices.\"\n                        f\" Got index of type {type(pidx.index)} at position {position}\")\n      else:\n        raise TypeError(f\"dynamic_slice: unrecognized index {pidx.index} at position {position}.\")\n\n    start_indices: list[ArrayLike] = []\n    slice_sizes: list[int] = []\n    rev_axes: list[int] = []\n    squeeze_axes: list[int] = []\n    newaxis_dims: list[int] = []\n\n    expanded = self.expand_ellipses()\n    trivial_slicing = True\n    for pidx in expanded.indices:\n      if pidx.typ in [IndexType.BOOLEAN, IndexType.ELLIPSIS]:\n        raise RuntimeError(f\"Internal: unexpected index encountered: {pidx}\")","sourceCodeStart":486,"sourceCodeEnd":522,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L486-L522","documentation":"In dynamic-slice lowering, indices classified as ARRAY type must be scalar (0-dimensional). If a sequence or non-scalar array is used as an index, JAX raises this TypeError because dynamic_slice only takes scalar start indices.","triggerScenarios":"Passing a list, tuple, or shape-(n,) array as an index in a dynamic index expression, e.g. x.at[[0, 1]].get() routed through to_dynamic_slice, or mixing gather-style arrays into dynamic slicing.","commonSituations":"Mixing NumPy fancy-indexing habits (arr[[i, j]]) with JAX dynamic slicing; passing an index array where a scalar tracer was expected in jit code; index variables accidentally wrapped in lists.","solutions":["Use standard fancy indexing (x[np.array([i, j])]) instead of the dynamic-slice path","Squeeze the index to a scalar: index = jnp.asarray(index).reshape(()) when it holds one element","Use lax.gather or x.at[indices_array].get() for multiple indices"],"exampleFix":"// before\nidx = [i]\ny = x.at[idx].get()  # sequence index\n// after\ny = x.at[i].get()  # scalar index","handlingStrategy":"type-guard","validationCode":"import numpy as np\nassert all(np.shape(i) == () and not isinstance(i, (list, tuple)) for i in idx_tuple if not isinstance(i, slice))","typeGuard":"def is_scalar_index(i) -> bool:\n    import numpy as np\n    return not isinstance(i, (list, tuple)) and np.shape(i) == ()","tryCatchPattern":"try:\n    y = x.at[idx].get()\nexcept TypeError as e:\n    if 'only scalar indices allowed' in str(e):\n        y = x[np.asarray(idx)]  # fall back to gather-style indexing\n    else:\n        raise","preventionTips":["Use scalar tracers or ints in dynamic-slice positions","Use array indexing (x[np.array([...])]) for multiple indices"],"tags":["jax","indexing","scalar-index","dynamic-slice"],"backgroundTag":"non-scalar-index-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}