Fix get vmap

PiperOrigin-RevId: 833888270

Author: justinjfu

jax/_src/state/primitives.py

@@ -835,6 +835,7 @@ def _batch_indexer(
           np.dtype('int32'), new_integer_indexer_shape, 0)
     else:
       batch_idx = indexing.Slice(0, axis_size)  # type: ignore
+      new_integer_indexer_shape = ()
     new_indices.insert(ref_dim, batch_idx)
   return indexing.NDIndexer(
       tuple(new_indices), ref_shape, new_integer_indexer_shape, validate=True
@@ -862,38 +863,41 @@ def _get_vmap(batched_args, batched_dims, *, tree):
   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 = bool(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)
   )
 
   out = get_p.bind(ref, *flat_indexers, tree=tree)
-  if not int_indexers_contiguous:  # will always be moved to the front
-    out_bdim = 0
-  else:  # originally not going to be moved to the front
-    if new_int_indexers_contiguous:  # now not going to be moved to the front
-      try:
-        out_bdim = is_new_int_indexing.index(True)
-      except ValueError:
-        out_bdim = 0
-    else:  # now going to be moved to the front
-      original_pos = is_int_indexing.index(True)
-      array_indexer_shape = new_indexers[0].int_indexer_shape
-      array_indexer_len = len(array_indexer_shape)
-
-      transpose_order = list(range(len(out.shape)))
-      transpose_order = (
-          transpose_order[0],
-          *transpose_order[array_indexer_len:array_indexer_len+original_pos],
-          *transpose_order[1:array_indexer_len],
-          *transpose_order[array_indexer_len+original_pos:],
-      )
+  if will_add_int_batcher and not new_int_indexers_contiguous:  # now going to be moved to the front
+    original_pos = is_int_indexing.index(True)
+    array_indexer_shape = new_indexers[0].int_indexer_shape
+    array_indexer_len = len(array_indexer_shape)
 
-      out = lax.transpose(out, transpose_order)
+    transpose_order = list(range(len(out.shape)))
+    transpose_order = (
+        transpose_order[0],
+        *transpose_order[array_indexer_len:array_indexer_len+original_pos],
+        *transpose_order[1:array_indexer_len],
+        *transpose_order[array_indexer_len+original_pos:],
+    )
+    out = lax.transpose(out, transpose_order)
+    out_bdim = 0
+  else:
+    if ref_dim is not batching.not_mapped:
+      # We only trigger this case when the int_indexer shape is empty,
+      # so we don't need to account for int_indexer_shape.
+      int_indexers_before_ref_dim = int(np.sum(is_new_int_indexing[:ref_dim]))
+      out_bdim = ref_dim - int_indexers_before_ref_dim
+    else:
       out_bdim = 0
+      if any(is_int_indexing):
+        original_pos = is_int_indexing.index(True)
+        out_bdim = original_pos
   return out, out_bdim
 batching.primitive_batchers[get_p] = _get_vmap
 

tests/state_test.py

@@ -136,6 +136,30 @@ def f(x_ref):
       self.assertEqual(out_aval.shape, out_shape)
       self.assertEqual(out_aval.dtype, out_dtype)
 
+  @parameterized.parameters(
+      ((4, 5), 0, (0,)),
+      ((4, 5), 1, (0,)),
+      ((9, 10, 11, 12), 0, (slice(None), 0, 1)),  # Contiguous int indexing
+      ((9, 10, 11, 12), 0, (0, slice(None), 1)),  # Non-contiguous int indexing
+      ((9, 10, 11, 12), 1, (slice(None), 0, 1)),  # Contiguous after batch
+      ((9, 10, 11, 12), 2, (slice(None), 0, 1)),  # Non-contiguous after batch
+      ((9, 10, 11, 12), 3, (slice(None), slice(None), 0)),
+      # Shaped int indexer, contiguous after batch
+      ((9, 10, 11, 12), 3,
+       (slice(None), slice(None), np.array([[0,1]]))),
+      # Shaped int indexer, non-contiguous after batch
+      ((9, 10, 11, 12), 2,
+       (np.array([[0, 1]]), slice(None), np.array([[0, 1]]))),
+  )
+  def test_vmap_of_get_regression(self, shape, in_axes, indexer):
+    # Regression test for https://github.com/jax-ml/jax/issues/33309
+    def f(x):
+      return x[indexer]
+    x = jnp.ones(shape)
+    result = jax.vmap(f, in_axes=in_axes)(jax.new_ref(x))
+    expected = jax.vmap(f, in_axes=in_axes)(x)
+    self.assertArraysEqual(result, expected)
+
   def test_swap_abstract_eval_must_take_in_refs(self):
     ref_aval = core.ShapedArray((), jnp.float32)
     val_aval = core.ShapedArray((), jnp.float32)