jax-ml/jax · error · NotImplementedError

No support for vmapping over nontrivial slices just yet: {id

Error message

No support for vmapping over nontrivial slices just yet: {idx}

What it means

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.

Source

Thrown at jax/_src/state/primitives.py:792

  new_indices: list[Array | indexing.Slice | int] = []
  new_integer_indexer_shape = (axis_size, *indexer.int_indexer_shape)
  for idx, dim in zip(indices, indices_dims):
    if idx_is_batched:
      # If at least one of the idx is batched, we broadcast them all and move the
      # batch dim to the front.
      if isinstance(idx, indexing.Slice):
        # size is static, but start can be dynamic
        # Check if start is static (which it can be)
        is_static_slice = len(tree_util.tree_leaves(idx)) == 0
        if is_static_slice:
          new_indices.append(idx)
          continue
        dim = dim.start
        if dim is None:
          # Broadcasting the slice is free (the start index stays the same)
          new_indices.append(idx)
        else:
          raise NotImplementedError(
              f"No support for vmapping over nontrivial slices just yet: {idx}")
      else:
        # Check if we are indexing with a scalar or not. If we are indexing
        # with a scalar and we are not batched, we can avoid broadcasting it.
        if not shapeof(idx):
          new_indices.append(idx)
        else:
          if dim is None:
            bcast_dims = tuple(range(1, np.ndim(idx) + 1))
            idx = lax.broadcast_in_dim(idx, new_integer_indexer_shape,
                                       bcast_dims)
          else:
            idx = batching.moveaxis(idx, dim, 0)
          new_indices.append(idx)
    else:
      if ref_dim is not None:
        if not isinstance(idx, indexing.Slice):
          if shapeof(idx):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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.
2. Hoist the varying offset out of vmap so slices have constant starts.
3. Unroll the batch with a Python loop or lax.map instead of vmap for this operation.

Example fix

// before
jax.vmap(lambda r, i: r.swap(slice(i, i+k), upd))(refs, starts)
// after
jax.vmap(lambda r, i: r.swap(jnp.arange(i, i+k), upd))(refs, starts)

def batched_slice_ok(idx) -> bool:
    return not isinstance(idx, slice) or idx.start is None

When it happens

Trigger: `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).

Common situations: 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.

Related errors

• ragged_dot vmap over any dim but 0 - NYI
• reduce_window batching is not implemented for initial values
• Batching over custom allocations is not supported yet.
• Batching over dynamic grid values is not supported yet.
• vmapping pallas_call with no arguments.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8754b10dc6878002. Report an issue: GitHub.