{"record":{"id":"8754b10dc6878002","repo":"jax-ml/jax","slug":"no-support-for-vmapping-over-nontrivial-slices-jus","errorCode":null,"errorMessage":"No support for vmapping over nontrivial slices just yet: {idx}","messagePattern":"No support for vmapping over nontrivial slices just yet: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":792,"sourceCode":"  new_indices: list[Array | indexing.Slice | int] = []\n  new_integer_indexer_shape = (axis_size, *indexer.int_indexer_shape)\n  for idx, dim in zip(indices, indices_dims):\n    if idx_is_batched:\n      # If at least one of the idx is batched, we broadcast them all and move the\n      # batch dim to the front.\n      if isinstance(idx, indexing.Slice):\n        # size is static, but start can be dynamic\n        # Check if start is static (which it can be)\n        is_static_slice = len(tree_util.tree_leaves(idx)) == 0\n        if is_static_slice:\n          new_indices.append(idx)\n          continue\n        dim = dim.start\n        if dim is None:\n          # Broadcasting the slice is free (the start index stays the same)\n          new_indices.append(idx)\n        else:\n          raise NotImplementedError(\n              f\"No support for vmapping over nontrivial slices just yet: {idx}\")\n      else:\n        # Check if we are indexing with a scalar or not. If we are indexing\n        # with a scalar and we are not batched, we can avoid broadcasting it.\n        if not shapeof(idx):\n          new_indices.append(idx)\n        else:\n          if dim is None:\n            bcast_dims = tuple(range(1, np.ndim(idx) + 1))\n            idx = lax.broadcast_in_dim(idx, new_integer_indexer_shape,\n                                       bcast_dims)\n          else:\n            idx = batching.moveaxis(idx, dim, 0)\n          new_indices.append(idx)\n    else:\n      if ref_dim is not None:\n        if not isinstance(idx, indexing.Slice):\n          if shapeof(idx):","sourceCodeStart":774,"sourceCodeEnd":810,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L774-L810","documentation":"The batching (vmap) rule for get/swap/addupdate supports slices only when the slice start is unaffected by batching. When a batched dimension flows into a slice whose start index must vary per-batch-element (dim.start is not None after unmapping), the required gather is not implemented and NotImplementedError is raised.","triggerScenarios":"`jax.vmap` over a function that indexes a Ref/getter with a slice like `ref[i:i+k]` where `i` is batched (depends on the vmapped argument).","commonSituations":"Vmapped rollout/attention windows that slice with per-example offsets; vmap of functions doing sliding-window updates on state buffers; migrating loops to vmap and hitting slice limits of the experimental state batching rules.","solutions":["Replace the slice with static bounds plus a gathered start: compute indices with `jnp.arange` and use integer indexing (gather) instead of a slice with a batched start.","Hoist the varying offset out of vmap so slices have constant starts.","Unroll the batch with a Python loop or lax.map instead of vmap for this operation."],"exampleFix":"// before\njax.vmap(lambda r, i: r.swap(slice(i, i+k), upd))(refs, starts)\n// after\njax.vmap(lambda r, i: r.swap(jnp.arange(i, i+k), upd))(refs, starts)","handlingStrategy":"fallback","validationCode":"def batched_slice_ok(idx) -> bool:\n    return not isinstance(idx, slice) or idx.start is None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer integer gather indices over slices with computed starts in vmap-able code.","Test state access under vmap early when adding batching.","Use lax.map as a fallback when vmap rules are missing."],"tags":["jax","vmap","batching","not-implemented","slicing"],"backgroundTag":"vmap-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}