jax-ml/jax · error · NotImplementedError

Batching with multiple indexers not supported.

Error message

Batching with multiple indexers not supported.

What it means

The vmap rule for `get` currently supports only a single indexer; if more than one index expression (e.g. a tuple of indices) is supplied while any indexer is batched, JAX raises NotImplementedError. This is a known TODO (multiple-indexer batching) in jax/_src/state/primitives.py.

Source

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

  )

def shapeof(x):
  return x.shape if isinstance(x, TransformedRef) else core.typeof(x).shape

def _get_vmap(batched_args, batched_dims, *, tree):
  axis_size, = {x.shape[d] for x, d in zip(batched_args, batched_dims)
                if d is not None}
  ref, *flat_idxs = batched_args
  ref_dim, *flat_idx_dims = batched_dims
  indexers = tree_util.tree_unflatten(tree, flat_idxs)
  if not indexers:
    return get_p.bind(ref, *flat_idxs, tree=tree), ref_dim
  indexers_dims = tree_util.tree_unflatten(tree, flat_idx_dims)

  idx_is_batched = any(i_dim is not None
                       for i_dim in flat_idx_dims)
  if len(indexers) > 1:
    raise NotImplementedError("Batching with multiple indexers not supported.")

  # TODO(sharadmv): handle vmap of multiple indexers
  new_indexers = tuple(_batch_indexer(indexer, dims, axis_size,
                                  ref.shape, ref_dim, idx_is_batched)
                     for indexer, dims in zip(indexers, indexers_dims))
  flat_indexers, tree = tree_util.tree_flatten(new_indexers)

  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexers[0])
  int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_int_indexing)[0]) == 1)
  )
  # Note: _batch_indexer will add a slice for the batch dim if the int_indexer
  # shape is empty, else it will use advanced/int indexing.
  will_add_int_batcher = ref_dim is not None and (idx_is_batched or indexers[0].int_indexer_shape)

  is_new_int_indexing, _, _ = indexing.unpack_ndindexer(new_indexers[0])
  new_int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_new_int_indexing)[0]) == 1)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten to a single integer index: linearize 2-D coordinates into one index, e.g. `i * ncols + j`, then index once.
2. Use one indexer plus slicing where possible, keeping the tuple length at 1.
3. Fall back to `lax.map` or a Python loop instead of vmap for multi-index access.

Example fix

// before
jax.vmap(lambda r, i, j: r.swap((i, j), v))(refs, rows, cols)
// after
flat = rows * ncols + cols
jax.vmap(lambda r, f: r.swap(f, v))(refs, flat)

# ensure a single indexer before vmap
assert len(indexers) == 1, 'multi-index get/swap not supported under vmap'
flat_idx = row * ncols + col  # linearize instead

When it happens

Trigger: `jax.vmap` over code calling `ref[i, j]` (tuple indexing → multiple indexers) where at least one index has a batched dim.

Common situations: Vmapped agents/models indexing 2-D state buffers with per-example (row, col) pairs; converting existing index-tuple code to batched execution with vmap.

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/88168825b1521798. Report an issue: GitHub.