Repository: jax-ml/jax
Pull request #23758 (2024-09-19 12:06:01)

[Pallas] Fix integer array indexing
[Pallas] Fix integer array indexing

Fixes https://github.com/google/jax/issues/22783

jax-fixit



Patch:
diff --git a/jax/_src/state/discharge.py b/jax/_src/state/discharge.py
index 7ab77d5b1..615fa862b 100644
--- a/jax/_src/state/discharge.py
+++ b/jax/_src/state/discharge.py
@@ -275,33 +275,97 @@ def _maybe_convert_to_dynamic_slice(
   return starts, sizes, squeeze_dims
 
 
-def _convert_to_array_indexer(indexer: indexing.NDIndexer
-                              ) -> tuple[int | Array, ...]:
-  # This is the general gather case. We need to create the gather arrays.
-  is_integer_indexer, _, integer_indexer = (
-      indexing.unpack_ndindexer(indexer)
+# In this code, indexing is handled in three ways: `slice`, `dynamic_slice`, and
+# gather. For the gather case, the goal is to create a gather array, which means
+# that we need to convert all other types of indexers into integer array
+# indexers. This is done by looping over all indexers and checking if they are
+# not integer array indexers, and if not, performing the conversion. However,
+# during this process, the indexing semantics may change. Specifically,
+# according to the indexing rules of NumPy, when there are integer array
+# indexers separated by other indexers, the axes corresponding to the integer
+# array indexers need to be moved to the front. After we convert all other
+# indexers to integer array indexers, the distinction between integer array
+# indexers and other types of indexers is lost. As a result, it becomes
+# impossible to determine which axes should be moved to the front. In this case,
+# we need to transpose the target array before the gather operation. We also
+# need to transpose the target array back after the gather operation, if it is
+# used in subsequent computations.
+def _maybe_transpose_before_gather(
+    indexer: indexing.NDIndexer
+) -> tuple[int, ...] | None:
+  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexer)
+
+  int_indexers_contiguous = bool(
+      np.all(np.diff(np.where(is_int_indexing)[0]) == 1)
   )
-  total_shape = indexer.get_indexer_shape()
-  int_indexer_shape = indexer.int_indexer_shape
-  slice_shape = total_shape[len(int_indexer_shape):]
-  slice_dims = tuple(
-      i + len(int_indexer_shape) for i in range(len(slice_shape))
+  if int_indexers_contiguous:
+    return None  # no transpose needed
+
+  int_indexer_idxs: list[int] = []
+  non_int_indexer_idxs: list[int] = []
+  for i, is_int_index in enumerate(is_int_indexing):
+    (int_indexer_idxs if is_int_index else non_int_indexer_idxs).append(i)
+  transpose_order = (*int_indexer_idxs, *non_int_indexer_idxs)
+  return transpose_order
+
+
+def _perform_transpose_before_gather(
+    target_arr: Array,
+    indexer: indexing.NDIndexer,
+    transpose_order: tuple[int, ...],
+) -> tuple[Array, indexing.NDIndexer]:
+  new_target_arr = target_arr.transpose(transpose_order)
+  reordered_indices = tuple(indexer.indices[i] for i in transpose_order)
+  new_indexer = indexing.NDIndexer(
+      indices=reordered_indices,
+      shape=indexer.shape,
+      int_indexer_shape=indexer.int_indexer_shape,
   )
-  slice_dim_iter = iter(slice_dims)
-  slice_indexer: list[Array] = []
-  for idx, is_int_index in zip(indexer.indices, is_integer_indexer):
-    if not is_int_index:
-      assert isinstance(idx, indexing.Slice)
-      slice_indices = lax.broadcasted_iota(
-          np.dtype("int32"), total_shape, next(slice_dim_iter)
-      ) * idx.stride + idx.start
-      slice_indexer.append(slice_indices)
-      integer_indexer = tuple(
-          lax.expand_dims(idx, (-1,)) for idx in integer_indexer
+  return new_target_arr, new_indexer
+
+
+def _convert_to_gather_arrays(indexer: indexing.NDIndexer) -> tuple[Array, ...]:
+  # This is the general gather case. We need to create the gather arrays.
+  total_shape = indexer.get_indexer_shape()
+  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexer)
+
+  if any(is_int_indexing):
+    n_idxers = len(indexer.indices)
+    int_indexer_shape = indexer.int_indexer_shape
+    n_int_indexers = sum(1 for p in is_int_indexing if p)
+    last_int_index_idx = n_idxers - 1 - is_int_indexing[::-1].index(True)
+    n_slice_index_dims_after_int = n_idxers - last_int_index_idx - 1
+
+  def get_idx_in_shape_after_indexing(i):
+    if not any(is_int_indexing):
+      return i
+
+    if i < n_idxers - n_slice_index_dims_after_int - n_int_indexers:
+      return i
+    if i < n_idxers - n_slice_index_dims_after_int:
+      raise ValueError
+    return i - n_int_indexers + len(int_indexer_shape)
+
+  arrs = []
+  for i, idxer in enumerate(indexer.indices):
+    if isinstance(idxer, indexing.Slice):
+      idx_in_shape_after_indexing = get_idx_in_shape_after_indexing(i)
+      arr = (
+          lax.iota(np.int32, total_shape[idx_in_shape_after_indexing])
+          * idxer.stride
+          + idxer.start
       )
-      continue
-  assert next(slice_dim_iter, None) is None
-  return tuple(merge_lists(is_integer_indexer, slice_indexer, integer_indexer))
+      diff = len(total_shape) - idx_in_shape_after_indexing - 1
+      arr = arr.reshape(arr.shape + (1,) * diff)
+      arrs.append(arr)
+    elif isinstance(idxer, (np.ndarray, Array)):
+      diff = n_idxers - 1 - last_int_index_idx
+      arr = idxer.reshape(idxer.shape + (1,) * diff)
+      arrs.append(arr)
+    else:
+      raise ValueError(f"Invalid type of idxer: {type(idxer).__name__}")
+
+  return tuple(arrs)
 
 
 @register_discharge_rule(get_p)
@@ -313,20 +377,8 @@ def _get_discharge_rule(
   y = _get_discharge(x, idx, tree)
   return (None,) * (len(idx) + 1), y
 
-def _prepend_gather(x, indexer):
-  # NumPy advanced int indexing won't prepend w/ only one dim, so add dummy.
-  return x[None][(np.array(0, 'int32'), *indexer)]
 
-def _prepend_scatter(x, indexer, val, *, add=False):
-  # NumPy advanced int indexing won't prepend w/ only one dim, so add dummy.
-  # However, since this is scatter, we need to remove the 1-sized dimension
-  # we added at the front.
-  if add:
-    return x[None].at[(0, *indexer)].add(val)[0]
-  return x[None].at[(0, *indexer)].set(val)[0]
-
-
-def _index_array(x, indexer):
+def _index_array(x, indexer: indexing.NDIndexer):
   if _is_trivial_indexer(indexer):
     return x
   # Try the three APIs in the following order: `lax.slice`,
@@ -336,13 +388,16 @@ def _index_array(x, indexer):
   # If everything in the indexer is a slice or ()-shaped, we can also
   # use `lax.dynamic_slice` with 1-sized slices for ()-shaped indices.
   # We need to squeeze out the 1-sized slices at the end.
-  elif maybe_slice := _maybe_convert_to_dynamic_slice(indexer):
-    starts, sizes, squeeze_dims = maybe_slice
+  elif maybe_dynamic_slice := _maybe_convert_to_dynamic_slice(indexer):
+    starts, sizes, squeeze_dims = maybe_dynamic_slice
     y = lax_slicing.dynamic_slice(x, starts, sizes)
     x = lax.squeeze(y, squeeze_dims)
   else:
-    indexer = _convert_to_array_indexer(indexer)
-    x = x[None][(np.array(0, "int32"), *indexer)]
+    transpose_order = _maybe_transpose_before_gather(indexer)
+    if transpose_order is not None:
+      x, indexer = _perform_transpose_before_gather(x, indexer, transpose_order)
+    arrays = _convert_to_gather_arrays(indexer)
+    x = x[arrays]
   return x
 
 
@@ -367,53 +422,79 @@ def transform_array(x, transforms):
 def transform_swap_array(x, transforms, val):
   if transforms is None:
     transforms = []
-  result = x
-  result_val = val
-  # Compute updated "val" (result).
-  _results = [x]
+
+  # Will hold the value read from `x` before the swap, and will have the same
+  # shape as `val`.
+  new_val = x
+  # List of intermediate results by transforming `x`.
+  intermediates = [x]
+
+  # Read phase (forward loop)
   for transform in transforms:
     match transform:
       case indexing.NDIndexer():
         indexer = transform
         if _is_trivial_indexer(indexer):
-          _results.append(_results[-1])
+          intermediates.append(intermediates[-1])
           continue
         # If everything in the indexer is a slice or ()-shaped, we can also
         # use `lax.dynamic_slice` with 1-sized slices for ()-shaped indices.
         # We need to squeeze out the 1-sized slices at the end.
         if maybe_slice := _maybe_convert_to_dynamic_slice(indexer):
           starts, sizes, squeeze_dims = maybe_slice
-          result_old = lax_slicing.dynamic_slice(result, starts, sizes)
-          result = lax.squeeze(result_old, squeeze_dims)
+          new_val = lax.squeeze(
+              lax_slicing.dynamic_slice(new_val, starts, sizes), squeeze_dims
+          )
         else:
-          indexer = _convert_to_array_indexer(indexer)
-          result = _prepend_gather(result, indexer)
-        _results.append(result)
+          transpose_order = _maybe_transpose_before_gather(indexer)
+          if transpose_order is not None:
+            new_val, indexer = _perform_transpose_before_gather(
+                new_val, indexer, transpose_order
+            )
+          arrays = _convert_to_gather_arrays(indexer)
+          new_val = new_val[arrays]
+          # Here, we don't need to transpose `new_val` back because it now holds
+          # the result of the indexing, and is no longer the original array that
+          # was indexed into.
+        intermediates.append(new_val)
       case RefBitcaster():
-        _results.append(bitcast(result, transform.dtype))
+        intermediates.append(bitcast(new_val, transform.dtype))
       case RefReshaper():
-        _results.append(result.reshape(transform.shape))
+        intermediates.append(new_val.reshape(transform.shape))
       case _:
         raise NotImplementedError(f"Unsupported transform: {transform}")
 
-  # Compute updated "x" (result_val)
-  for i, transform in reversed(list(enumerate(transforms))):
+  # Will hold the final state of the `x` after `val` has been written to the
+  # transformed location, and will have the same shape as `x`.
+  new_x = val
+
+  # Write phase (reversed loop)
+  for intermediate, transform in reversed(zip(intermediates[:-1], transforms)):
     if isinstance(transform, indexing.NDIndexer):
       indexer = transform
       if _is_trivial_indexer(indexer):
         continue
       if maybe_slice := _maybe_convert_to_dynamic_slice(indexer):
         starts, _, squeeze_dims = maybe_slice
-        result_val = lax.expand_dims(result_val, squeeze_dims)
-        result_val = lax_slicing.dynamic_update_slice(
-            _results[i], result_val, starts
+        new_x = lax_slicing.dynamic_update_slice(
+            intermediate, lax.expand_dims(new_x, squeeze_dims), starts
         )
       else:
-        indexer = _convert_to_array_indexer(indexer)
-        result_val = _prepend_scatter(_results[i], indexer, result_val)
+        transpose_order = _maybe_transpose_before_gather(indexer)
+        if transpose_order is not None:
+          intermediate, indexer = _perform_transpose_before_gather(
+              intermediate, indexer, transpose_order
+          )
+        arrays = _convert_to_gather_arrays(indexer)
+        new_x = intermediate.at[arrays].set(new_x)  # pytype: disable=attribute-error
+        if transpose_order is not None:
+          transpose_order_inversed = np.argsort(transpose_order)
+          new_x = new_x.transpose(transpose_order_inversed)
     else:
       raise NotImplementedError(f"Unsupported transform: {transform}")
-  return result, result_val
+
+  return new_val, new_x
+
 
 def _get_discharge(x, idx, tree):
   transforms = tree_util.tree_unflatten(tree, idx)
@@ -446,8 +527,10 @@ def _addupdate_discharge(x, val, idx, tree):
   if len(transforms) > 1:
     raise NotImplementedError("Only single indexer is supported.")
   indexer = transforms[0]
+
   if _is_trivial_indexer(indexer):
     return x + val
+
   # If everything in the indexer is a slice or ()-shaped, we can also
   # use `lax.dynamic_slice` with 1-sized slices for ()-shaped indices.
   # We need to squeeze out the 1-sized slices at the end.
@@ -457,8 +540,17 @@ def _addupdate_discharge(x, val, idx, tree):
     val = lax.expand_dims(val, squeeze_dims)
     y = lax_slicing.dynamic_update_slice(x, x_old + val, starts)
     return y
-  indexer = _convert_to_array_indexer(indexer)
-  return _prepend_scatter(x, indexer, val, add=True)
+
+  transpose_order = _maybe_transpose_before_gather(indexer)
+  if transpose_order is not None:
+    x, indexer = _perform_transpose_before_gather(x, indexer, transpose_order)
+  arrays = _convert_to_gather_arrays(indexer)
+  x = x.at[arrays].add(val)
+  if transpose_order is not None:
+    transpose_order_inversed = np.argsort(transpose_order)
+    x = x.transpose(transpose_order_inversed)
+  return x
+
 
 @weakref_lru_cache
 def _cached_closed_jaxpr_discharge(closed_jaxpr: core.ClosedJaxpr):
diff --git a/jax/_src/state/indexing.py b/jax/_src/state/indexing.py
index e7b581680..e6e6b8a5e 100644
--- a/jax/_src/state/indexing.py
+++ b/jax/_src/state/indexing.py
@@ -272,11 +272,21 @@ class NDIndexer:
     return cls(indices, shape, int_indexer_shape, validate=True)
 
   def get_indexer_shape(self) -> tuple[int | Array, ...]:
-    _, slice_indexers, _ = unpack_ndindexer(self)
-    slice_shape = [s.size for s in slice_indexers]
-    # In NDIndexers, the int_indexer_shape is *always* at the front of the
-    # result.
-    return (*self.int_indexer_shape, *slice_shape)
+    is_int_indexing, slice_indexers, _ = unpack_ndindexer(self)
+
+    slice_shape = tuple(s.size for s in slice_indexers)
+    int_indexers_contiguous = bool(
+        np.all(np.diff(np.where(is_int_indexing)[0]) == 1)
+    )
+    if not int_indexers_contiguous:
+      return self.int_indexer_shape + slice_shape
+
+    has_int_indexers = any(is_int_indexing)
+    if has_int_indexers:
+      pos = is_int_indexing.index(True)
+      return slice_shape[:pos] + self.int_indexer_shape + slice_shape[pos:]
+
+    return slice_shape
 
   def transform_shape(self, shape: None | tuple[int | Array, ...]) -> None | tuple[int | Array, ...]:
     del shape  # Unused
diff --git a/jax/_src/state/primitives.py b/jax/_src/state/primitives.py
index 578ef0bbc..1237da57f 100644
--- a/jax/_src/state/primitives.py
+++ b/jax/_src/state/primitives.py
@@ -454,11 +454,52 @@ pe.partial_eval_jaxpr_custom_rules[addupdate_p] = partial(
 
 ##  get/swap/addupdate batching rules
 
-def _batch_indexer(indexer: indexing.NDIndexer, dims,
-                   axis_size: int,
-                   ref_shape: tuple[int, ...],
-                   ref_dim: int | batching.NotMapped,
-                   idx_is_batched: bool) -> indexing.NDIndexer:
+def _batch_indexer(
+    indexer: indexing.NDIndexer,
+    dims,
+    axis_size: int,
+    ref_shape: tuple[int, ...],
+    ref_dim: int | batching.NotMapped,
+    idx_is_batched: bool,
+) -> indexing.NDIndexer:
+  """Converts a batched indexer into an unbatched one.
+
+  This function handles the complexity of `vmap`-style batching where either the
+  `ref` being indexed, the indexer, or both may have batched dimensions. The
+  goal is to produce a new indexer that acts as if applied in a batched context,
+  but without actual batching, enabling downstream code to process it as usual.
+
+  If any index in `indexer` is batched, all array indexers are normalized. If
+  the array indexer contains a batched dimension, the dimension is moved to the
+  front (axis 0). If the array indexer not batched, it is broadcasted to include
+  a batch dimension at the front. This is to guarantee that all array indexers
+  are still of the same shape.
+
+  Slices are passed through unchanged unless they contain dynamic elements and
+  are themselves batched, which is currently unsupported.
+
+  If `ref` is batched (`ref_dim` is not `NotMapped`), we simulate per-example
+  indexing by inserting a new iota array at the position corresponding to
+  `ref_dim` in the indexer.
+
+  It is worth noting that if the array indexers in the original indexer are
+  contiguous, but become non-contiguous in the new indexer due to the insertion
+  of the iota, the dimensions corresponding to the array indexers will be moved
+  to the front in the indexing result. The batched dimension will be at axis 0,
+  while the dimensions corresponding to the array indexers in the original
+  indexer will start from axis 1. This behavior would cause a mismatch between
+  the original indexer and the new indexer. Callers must take this behavior into
+  account and properly transpose the arrays involved to avoid this mismatch.
+
+  Args:
+    indexer: An `NDIndexer` that indexes into `ref`.
+    dims: A pytree with the same structure as `indexer`, indicating which
+      dimension (if any) is batched for each array indexer.
+    axis_size: Size of the batch dimension.
+    ref_shape: Shape of `ref`.
+    ref_dim: The dimension of `ref` that is batched (if any).
+    idx_is_batched: Whether any index in the `indexer` is batched.
+  """
   indices = indexer.indices
   indices_dims = dims.indices
   new_indices: list[Array | indexing.Slice | int] = []
@@ -510,9 +551,9 @@ def _batch_indexer(indexer: indexing.NDIndexer, dims,
   if ref_dim is not batching.not_mapped:
     iota = lax.broadcasted_iota(np.dtype('int32'), new_integer_indexer_shape, 0)
     new_indices.insert(ref_dim, iota)
-  return indexing.NDIndexer(tuple(new_indices), ref_shape,
-                            new_integer_indexer_shape,
-                            validate=True)
+  return indexing.NDIndexer(
+      tuple(new_indices), ref_shape, new_integer_indexer_shape, validate=True
+  )
 
 def _get_vmap(batched_args, batched_dims, *, tree):
   axis_size, = {x.shape[d] for x, d in zip(batched_args, batched_dims)
@@ -527,11 +568,42 @@ def _get_vmap(batched_args, batched_dims, *, tree):
   if len(indexers) > 1:
     raise NotImplementedError("Batching with multiple indexers not supported.")
   # TODO(sharadmv): handle vmap of multiple indexers
-  indexers = tuple(_batch_indexer(indexer, dims, axis_size,
+  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(indexers)
-  return get_p.bind(ref, *flat_indexers, tree=tree), 0
+  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)
+  )
+  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
+      out_bdim = is_new_int_indexing.index(True)
+    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:],
+      )
+
+      out = lax.transpose(out, transpose_order)
+      out_bdim = 0
+  return out, out_bdim
 batching.primitive_batchers[get_p] = _get_vmap
 
 def _swap_vmap(batched_args, batched_dims, *, tree):
@@ -549,15 +621,59 @@ def _swap_vmap(batched_args, batched_dims, *, tree):
   if len(indexers) > 1:
     raise NotImplementedError("Batching with multiple indexers not supported.")
   # TODO(sharadmv): handle vmap of multiple indexers
-  indexers = tuple(_batch_indexer(indexer, dims, axis_size,
+  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(indexers)
-  if (ref_is_batched or idx_is_batched) and not val_is_batched:
-    val = batching.broadcast(val, axis_size, 0)
-  if val_is_batched:
-    val = batching.moveaxis(val, val_dim, 0)
-  return swap_p.bind(ref, val, *flat_indexers, tree=tree), 0
+  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)
+  )
+  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)
+  )
+
+  if not new_int_indexers_contiguous:  # will be moved to the front
+    batched_dim_in_result = 0
+  else:
+    batched_dim_in_result = is_new_int_indexing.index(True) + 0
+
+  if not val_is_batched:
+    if ref_is_batched or idx_is_batched:
+      val = batching.broadcast(val, axis_size, batched_dim_in_result)
+  else:
+    val = batching.moveaxis(val, val_dim, batched_dim_in_result)
+
+  transpose_order_inversed = None
+
+  # Originally not going to be moved to the front, but now going to be moved to
+  # the front.
+  if int_indexers_contiguous and not new_int_indexers_contiguous:
+    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(val.shape)))
+    transpose_order = (
+        transpose_order[0],
+        *transpose_order[1+original_pos:(1+original_pos)+(array_indexer_len-1)],
+        *transpose_order[1:1+original_pos],
+        *transpose_order[(1+original_pos)+(array_indexer_len-1):],
+    )
+    val = val.transpose(transpose_order)
+    transpose_order_inversed = np.argsort(transpose_order)
+
+  out = swap_p.bind(ref, val, *flat_indexers, tree=tree)
+
+  # `val` should not be transposed, but we needed to transpose it to match
+  # `swap_p`. As a result, the output of `swap_p` is also transposed. Now we
+  # need to transpose it back.
+  if transpose_order_inversed is not None:
+    out = out.transpose(transpose_order_inversed)
+
+  return out, batched_dim_in_result
 batching.primitive_batchers[swap_p] = _swap_vmap
 
 def _addupdate_vmap(batched_args, batched_dims, *, tree):
@@ -575,14 +691,47 @@ def _addupdate_vmap(batched_args, batched_dims, *, tree):
   if len(indexers) > 1:
     raise NotImplementedError("Batching with multiple indexers not supported.")
   # TODO(sharadmv): handle vmap of multiple indexers
-  indexers = tuple(_batch_indexer(indexer, dims, axis_size,
+  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(indexers)
-  if (ref_is_batched or idx_is_batched) and not val_is_batched:
-    val = batching.broadcast(val, axis_size, 0)
-  if val_is_batched:
-    val = batching.moveaxis(val, val_dim, 0)
+  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)
+  )
+  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)
+  )
+
+  if not new_int_indexers_contiguous:  # will be moved to the front
+    batched_dim_in_result = 0
+  else:
+    batched_dim_in_result = is_new_int_indexing.index(True)
+
+  if not val_is_batched:
+    if ref_is_batched or idx_is_batched:
+      val = batching.broadcast(val, axis_size, batched_dim_in_result)
+  else:
+    val = batching.moveaxis(val, val_dim, batched_dim_in_result)
+
+  # Originally not going to be moved to the front, but now going to be moved to
+  # the front.
+  if int_indexers_contiguous and not new_int_indexers_contiguous:
+    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(val.shape)))
+    transpose_order = (
+        transpose_order[0],
+        *transpose_order[1+original_pos:(1+original_pos)+(array_indexer_len-1)],
+        *transpose_order[1:1+original_pos],
+        *transpose_order[(1+original_pos)+(array_indexer_len-1):],
+    )
+    val = val.transpose(transpose_order)
+
   return addupdate_p.bind(ref, val, *flat_indexers, tree=tree), []
 batching.primitive_batchers[addupdate_p] = _addupdate_vmap
 


Test patch:
diff --git a/tests/pallas/indexing_test.py b/tests/pallas/indexing_test.py
index 6e9d552e3..3de0c1f30 100644
--- a/tests/pallas/indexing_test.py
+++ b/tests/pallas/indexing_test.py
@@ -218,12 +218,13 @@ class IndexerTest(jtu.JaxTestCase):
 
     indices = (ds(0, 2), np.arange(5)[:, None], np.arange(4)[None])
     indexer = NDIndexer.from_indices_shape(indices, shape)
-    self.assertTupleEqual(indexer.get_indexer_shape(), (5, 4, 2))
+    self.assertTupleEqual(indexer.get_indexer_shape(), (2, 5, 4))
 
   @hp.given(hps.data())
   def test_ndindexer(self, data):
     shape = data.draw(hnp.array_shapes())
     indexer = data.draw(nd_indexer_strategy(shape))
+
     is_int_indexer = [not isinstance(idx, Slice) for idx in indexer.indices]
     rest_indexers, int_indexers = util.partition_list(
         is_int_indexer, indexer.indices
@@ -235,16 +236,12 @@ class IndexerTest(jtu.JaxTestCase):
     self.assertTupleEqual(
         indexer.int_indexer_shape, expected_int_indexer_shape
     )
+
     for idx in rest_indexers:
       self.assertIsInstance(idx, (np.ndarray, Slice))
       if isinstance(idx, np.ndarray):
         self.assertTupleEqual(idx.shape, ())
         self.assertEqual(idx.dtype, np.dtype("int32"))
-    rest_shape = tuple(
-        r.size for r in rest_indexers if not isinstance(r, np.ndarray)
-    )
-    self.assertTupleEqual((*indexer.int_indexer_shape, *rest_shape),
-                          indexer.get_indexer_shape())
 
 
 @jtu.thread_unsafe_test_class()  # hypothesis is not thread safe
@@ -692,18 +689,18 @@ _ADVANCED_INDEXER_TEST_CASES = [
     ((4, 3), lambda arr, a, b, c, d: arr[a, 2]),
     # slice + 1-D array
     ((4, 3), lambda arr, a, b, c, d: arr[a, :]),
-    # ((4, 3), lambda arr, a, b, c, d: arr[:, a]),
+    ((4, 3), lambda arr, a, b, c, d: arr[:, a]),
     ((6, 8, 3), lambda arr, a, b, c, d: arr[c, ::3]),
-    # ((8, 6, 3), lambda arr, a, b, c, d: arr[::3, c]),
-    # ((8, 8, 3), lambda arr, a, b, c, d: arr[::4, ::2, a]),
-    # ((8, 8, 3), lambda arr, a, b, c, d: arr[::4, a, ::2]),
+    ((8, 6, 3), lambda arr, a, b, c, d: arr[::3, c]),
+    ((8, 8, 3), lambda arr, a, b, c, d: arr[::4, ::2, a]),
+    ((8, 8, 3), lambda arr, a, b, c, d: arr[::4, a, ::2]),
     ((8, 8, 3, 7), lambda arr, a, b, c, d: arr[b, ::4, a, ::2]),
     ((3, 8, 8, 7), lambda arr, a, b, c, d: arr[b, a, ::4, ::2]),
     # ((8, 8, 3, 7), lambda arr, a, b, c, d: arr[::4, b, a, ::2]),
     ((16, 3, 6, 2), lambda arr, a, b, c, d: arr[::4, a, 1::2, b]),
     ((8, 8, 3, 6), lambda arr, a, b, c, d: arr[b, ::4, a, a]),
     # slice + array w/ broadcasting
-    ((8, 8, 3, 6), lambda arr, a, b, c, d: \
+    ((8, 8, 3, 6), lambda arr, a, b, c, d:
         arr[b[:, None], ::4, a[None], a[:, None]]),
     # integer + slice + 1-D array
     ((5, 8, 8, 3), lambda arr, a, b, c, d: arr[2, ::4, ::2, a]),
diff --git a/tests/state_test.py b/tests/state_test.py
index ab11ab829..65f6f0427 100644
--- a/tests/state_test.py
+++ b/tests/state_test.py
@@ -28,6 +28,7 @@ from jax import random
 from jax import lax
 from jax._src import core
 from jax._src import config
+from jax._src import dtypes
 from jax._src import linear_util as lu
 from jax._src.interpreters import partial_eval as pe
 from jax._src import test_util as jtu
@@ -477,27 +478,17 @@ class StatePrimitivesTest(jtu.JaxTestCase):
     op=[
         lambda x_ref, indexer: [x_ref[indexer]],
         lambda x_ref, indexer: [
-            ref_swap(x_ref, indexer,
-                            jnp.ones(x_ref.shape, x_ref.dtype)[None][(0,
-                              *indexer)])],
+            ref_swap(x_ref, indexer, jnp.ones_like(x_ref[indexer]))],
         lambda x_ref, indexer: (
-            ref_addupdate(x_ref, indexer,
-                                jnp.ones(x_ref.shape, x_ref.dtype)[None][(0,
-                                  *indexer)])
-            or [jnp.ones(x_ref.shape, x_ref.dtype)[None][(0, *indexer)]])
+            ref_addupdate(x_ref, indexer, jnp.ones_like(x_ref[indexer]))
+            or [jnp.ones_like(x_ref[indexer])]),
     ],
   )
   def test_vmap(self, ref_shape, ref_bdim, idx_shape, indexed_dims,
                     idx_bdims, out_bdim, op):
-
-    float_ = (jnp.dtype('float64') if config.enable_x64.value else
-              jnp.dtype('float32'))
-    int_ = (jnp.dtype('int64') if config.enable_x64.value else
-            jnp.dtype('int32'))
+    intx = dtypes.canonicalize_dtype(jnp.int64)
+    floatx = dtypes.canonicalize_dtype(jnp.float64)
     axis_size = 7
-    out_shape = tuple(d for d, b in zip(ref_shape, indexed_dims) if not b)
-    if any(indexed_dims):
-      out_shape = (*idx_shape, *out_shape)
 
     def maybe_insert(shape, idx):
       if idx is None:
@@ -505,13 +496,13 @@ class StatePrimitivesTest(jtu.JaxTestCase):
       return tuple_insert(shape, idx, axis_size)
 
     batched_ref_shape = maybe_insert(ref_shape, ref_bdim)
-    ref_aval = shaped_array_ref(ref_shape, float_)
-    bat_ref_aval = shaped_array_ref(batched_ref_shape, float_)
+    ref_aval = shaped_array_ref(ref_shape, floatx)
+    bat_ref_aval = shaped_array_ref(batched_ref_shape, floatx)
 
-    idx_avals = [core.ShapedArray(idx_shape, int_)
+    idx_avals = [core.ShapedArray(idx_shape, intx)
                  for _ in idx_bdims]
     bat_idx_avals = [
-        core.ShapedArray(maybe_insert(idx_shape, idx_bdim), int_)
+        core.ShapedArray(maybe_insert(idx_shape, idx_bdim), intx)
         for idx_bdim in idx_bdims]
 
     def f(x_ref, *idxs):
@@ -531,6 +522,7 @@ class StatePrimitivesTest(jtu.JaxTestCase):
         wrap_init(f_batched, 1 + len(bat_idx_avals)), [bat_ref_aval, *bat_idx_avals])
     jaxpr, consts = discharge_state(stateful_jaxpr, stateful_consts)
     discharge_of_vmap_ans = core.eval_jaxpr(jaxpr, consts, a, *idxs)
+
     # vmap-of-discharge
     stateful_jaxpr, _, stateful_consts, () = pe.trace_to_jaxpr_dynamic(
         wrap_init(f, 1 + len(idx_avals)), [ref_aval, *idx_avals])
@@ -830,11 +822,22 @@ if CAN_USE_HYPOTHESIS:
                                   min_size=len(ref_shape),
                                   max_size=len(ref_shape)))
     idx_shape = draw(hnp.array_shapes(max_dims=3, max_side=5))
-    if any(indexed_dims):
-      sliced_shape = (s for s, b in zip(ref_shape, indexed_dims) if not b)
-      slice_shape = (*idx_shape, *sliced_shape)
-    else:
+    if not any(indexed_dims):
       slice_shape = ref_shape
+    else:
+      sliced_shape = tuple(s for s, b in zip(ref_shape, indexed_dims) if not b)
+      int_indexers_contiguous = bool(
+          np.all(np.diff(np.where(indexed_dims)[0]) == 1)
+      )
+      if not int_indexers_contiguous:
+        slice_shape = (*idx_shape, *sliced_shape)
+      else:
+        insert_pos = indexed_dims.index(True)
+        slice_shape = (
+          *sliced_shape[:insert_pos],
+          *idx_shape,
+          *sliced_shape[insert_pos:],
+        )
     ref_aval = shaped_array_ref(ref_shape, np.float32)
     idx_avals = tuple(core.ShapedArray(idx_shape, np.int32) for _ in
         range(sum(indexed_dims)))