{"record":{"id":"88168825b1521798","repo":"jax-ml/jax","slug":"batching-with-multiple-indexers-not-supported","errorCode":null,"errorMessage":"Batching with multiple indexers not supported.","messagePattern":"Batching with multiple indexers not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":848,"sourceCode":"  )\n\ndef shapeof(x):\n  return x.shape if isinstance(x, TransformedRef) else core.typeof(x).shape\n\ndef _get_vmap(batched_args, batched_dims, *, tree):\n  axis_size, = {x.shape[d] for x, d in zip(batched_args, batched_dims)\n                if d is not None}\n  ref, *flat_idxs = batched_args\n  ref_dim, *flat_idx_dims = batched_dims\n  indexers = tree_util.tree_unflatten(tree, flat_idxs)\n  if not indexers:\n    return get_p.bind(ref, *flat_idxs, tree=tree), ref_dim\n  indexers_dims = tree_util.tree_unflatten(tree, flat_idx_dims)\n\n  idx_is_batched = any(i_dim is not None\n                       for i_dim in flat_idx_dims)\n  if len(indexers) > 1:\n    raise NotImplementedError(\"Batching with multiple indexers not supported.\")\n\n  # TODO(sharadmv): handle vmap of multiple indexers\n  new_indexers = tuple(_batch_indexer(indexer, dims, axis_size,\n                                  ref.shape, ref_dim, idx_is_batched)\n                     for indexer, dims in zip(indexers, indexers_dims))\n  flat_indexers, tree = tree_util.tree_flatten(new_indexers)\n\n  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexers[0])\n  int_indexers_contiguous = bool(\n      np.all(np.diff(np.where(is_int_indexing)[0]) == 1)\n  )\n  # Note: _batch_indexer will add a slice for the batch dim if the int_indexer\n  # shape is empty, else it will use advanced/int indexing.\n  will_add_int_batcher = ref_dim is not None and (idx_is_batched or indexers[0].int_indexer_shape)\n\n  is_new_int_indexing, _, _ = indexing.unpack_ndindexer(new_indexers[0])\n  new_int_indexers_contiguous = bool(\n      np.all(np.diff(np.where(is_new_int_indexing)[0]) == 1)","sourceCodeStart":830,"sourceCodeEnd":866,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L830-L866","documentation":"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.","triggerScenarios":"`jax.vmap` over code calling `ref[i, j]` (tuple indexing → multiple indexers) where at least one index has a batched dim.","commonSituations":"Vmapped agents/models indexing 2-D state buffers with per-example (row, col) pairs; converting existing index-tuple code to batched execution with vmap.","solutions":["Flatten to a single integer index: linearize 2-D coordinates into one index, e.g. `i * ncols + j`, then index once.","Use one indexer plus slicing where possible, keeping the tuple length at 1.","Fall back to `lax.map` or a Python loop instead of vmap for multi-index access."],"exampleFix":"// before\njax.vmap(lambda r, i, j: r.swap((i, j), v))(refs, rows, cols)\n// after\nflat = rows * ncols + cols\njax.vmap(lambda r, f: r.swap(f, v))(refs, flat)","handlingStrategy":"fallback","validationCode":"# ensure a single indexer before vmap\nassert len(indexers) == 1, 'multi-index get/swap not supported under vmap'\nflat_idx = row * ncols + col  # linearize instead","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Linearize multi-dimensional indices into one flat index.","Keep indexer tuples length-1 in code paths you plan to vmap.","Check jax release notes; multiple-indexer batching support may land later."],"tags":["jax","vmap","batching","not-implemented","multiple-indexing"],"backgroundTag":"vmap-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}