Repository: google/jax

Issue:
Different result while using ``at[].apply`` inside jax.vmap
### Description

Jax's Array index helper `at` potentially has a bug when used inside `jax.vmap`. Using ``apply(lambda...)``  inside ``jax.vmap`` gives me a different result as to when run sequentially and inside jit. 
Here's the following code that I'm trying to run
```python
import jax
import jax.numpy as np

a = np.array([1,2,3,4,5,6])
def adder(a,val):
    a = a.at[val].apply(lambda a:a-1)
    return a

ind = np.array([3])
print(adder(a,ind))
```
I get the expected result;
``[1 2 3 4 5 6]``
However, when the same function is used inside jax.vmap I get a different result
```python
v_adder = jax.vmap(adder,in_axes=(0,None))
a = np.array([[1,2,3,4,5,6]])
ind = np.array([3])
print(v_adder(a,ind))
```
``[[1 2 3 0 5 6]]``
Note that I had to add another dimension in order to use it inside vmap which I don't think is the root cause.
Took me a day of debugging to hunt down this bug 😭. 


### What jax/jaxlib version are you using?

jax==0.4.13 jaxlib==0.4.13

### Which accelerator(s) are you using?

CPU

### Additional system info

WSL

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report, and sorry for the difficulty in tracking down the bug. We really appreciate it though, especially with the minimal repro.

> I get the expected result;
> [1 2 3 4 5 6]

Did you mean `[1 2 3 3 5 6]`?

This certainly looks like a bug to me! It reproduces on TPU too. Here are the jaxprs:

```python
print(jax.make_jaxpr(adder)(a,ind))
```
```
{ lambda ; a:i32[6] b:i32[1]. let
    c:bool[1] = lt b 0
    d:i32[1] = add b 6
    e:i32[1] = select_n c b d
    f:i32[1,1] = broadcast_in_dim[broadcast_dimensions=(0,) shape=(1, 1)] e
    _:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    g:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    h:i32[6] = scatter[
      dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,))
      indices_are_sorted=False
      mode=GatherScatterMode.FILL_OR_DROP
      unique_indices=False
      update_consts=()
      update_jaxpr={ lambda ; i:i32[] j:i32[]. let k:i32[] = sub i 1 in (k,) }
    ] a f g
  in (h,) }
```

```python
print(jax.make_jaxpr(v_adder)(a,ind))
```

```
{ lambda ; a:i32[1,6] b:i32[1]. let
    c:bool[1] = lt b 0
    d:i32[1] = add b 6
    e:i32[1] = select_n c b d
    f:i32[1,1] = broadcast_in_dim[broadcast_dimensions=(0,) shape=(1, 1)] e
    _:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    g:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    h:i32[1,1] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 1)] g
    i:i32[1,6] = scatter[
      dimension_numbers=ScatterDimensionNumbers(update_window_dims=(0,), inserted_window_dims=(1,), scatter_dims_to_operand_dims=(1,))
      indices_are_sorted=False
      mode=GatherScatterMode.FILL_OR_DROP
      unique_indices=False
      update_consts=()
      update_jaxpr={ lambda ; j:i32[] k:i32[]. let  in (k,) }
    ] a f h
  in (i,) }
```

What's up with that second `update_jaxpr`!?
Yeah sorry, scatter-apply just isn't batchable; there's a rule defined but it doesn't work. In particular, it [ignores the `update_jaxpr` (and hence the user-defined fucntion)](https://github.com/google/jax/blob/1795b12a9f7bca118d75c981fadd61331680de57/jax/_src/lax/slicing.py#L2070) it's given.

Resolving patch:
diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -1383,11 +1383,12 @@ def _dynamic_update_slice_batching_rule(batched_args, batch_dims):
                                   inserted_window_dims=(),
                                   scatter_dims_to_operand_dims=dims)
   index, index_bdim = _batch_dynamic_slice_indices(start_idx, start_idx_bd)
-  return _scatter_batching_rule(
-    scatter, (operand, index, update), (operand_bd, index_bdim, update_bd),
-    update_jaxpr=None, update_consts=None, dimension_numbers=dnums,
-    indices_are_sorted=True, unique_indices=True,
-    mode=GatherScatterMode.CLIP)
+  return jax.vmap(
+    partial(scatter, dimension_numbers=dnums,
+            indices_are_sorted=True, unique_indices=True,
+            mode=GatherScatterMode.CLIP),
+    in_axes=(operand_bd, index_bdim, update_bd),
+    out_axes=0)(operand, index, update), 0
 
 
 dynamic_update_slice_p = standard_primitive(
@@ -2067,7 +2068,6 @@ def _scatter_batching_rule(scatter_op, batched_args, batch_dims, *,
                            indices_are_sorted, unique_indices, mode):
   operand, indices, updates = batched_args
   operand_bdim, indices_bdim, updates_bdim = batch_dims
-  del update_jaxpr, update_consts  # Unused.
 
   # move the operand batch dim to the front if it is not None, otherwise create
   # it at the front (so that we can scatter into it)
@@ -2086,10 +2086,10 @@ def _scatter_batching_rule(scatter_op, batched_args, batch_dims, *,
         update_window_dims=update_window_dims,
         inserted_window_dims=inserted_window_dims,
         scatter_dims_to_operand_dims=scatter_dims_to_operand_dims)
-    return scatter_op(
-      operand, indices, updates, dnums,
+    return scatter_op.bind(
+      operand, indices, updates, dimension_numbers=dnums,
       indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
-      mode=mode), 0
+      mode=mode, update_jaxpr=update_jaxpr, update_consts=update_consts), 0
 
 
   # see the third case in _gather_batching_rule for comparison and comments
@@ -2108,10 +2108,10 @@ def _scatter_batching_rule(scatter_op, batched_args, batch_dims, *,
       update_window_dims=update_window_dims,
       inserted_window_dims=inserted_window_dims,
       scatter_dims_to_operand_dims=scatter_dims_to_operand_dims)
-  return scatter_op(
-      operand, indices, updates, dnums,
+  return scatter_op.bind(
+      operand, indices, updates, dimension_numbers=dnums,
       indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
-      mode=mode), 0
+      mode=mode, update_jaxpr=update_jaxpr, update_consts=update_consts), 0
 
 scatter_add_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-add',
@@ -2119,7 +2119,7 @@ def _scatter_batching_rule(scatter_op, batched_args, batch_dims, *,
 ad.primitive_jvps[scatter_add_p] = _scatter_add_jvp
 ad.primitive_transposes[scatter_add_p] = _scatter_add_transpose_rule
 batching.primitive_batchers[scatter_add_p] = (
-  partial(_scatter_batching_rule, scatter_add))
+  partial(_scatter_batching_rule, scatter_add_p))
 
 scatter_mul_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-mul',
@@ -2141,7 +2141,7 @@ def _scatter_mul_jvp_rhs(g, x, i, y, *, dimension_numbers,
           _scatter_mul_jvp_rhs)
 ad.primitive_transposes[scatter_mul_p] = _scatter_mul_transpose_rule
 batching.primitive_batchers[scatter_mul_p] = (
-  partial(_scatter_batching_rule, scatter_mul))
+  partial(_scatter_batching_rule, scatter_mul_p))
 
 def _scatter_extremal_jvp(scatter_op, primals, tangents, update_jaxpr,
                           update_consts, dimension_numbers,
@@ -2248,14 +2248,14 @@ def _scatter_extremal_jvp(scatter_op, primals, tangents, update_jaxpr,
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-min',
     weak_type_rule=_argnum_weak_type(0))
 batching.primitive_batchers[scatter_min_p] = (
-  partial(_scatter_batching_rule, scatter_min))
+  partial(_scatter_batching_rule, scatter_min_p))
 ad.primitive_jvps[scatter_min_p] = partial(_scatter_extremal_jvp, scatter_min_p)
 
 scatter_max_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-max',
     weak_type_rule=_argnum_weak_type(0))
 batching.primitive_batchers[scatter_max_p] = (
-  partial(_scatter_batching_rule, scatter_max))
+  partial(_scatter_batching_rule, scatter_max_p))
 ad.primitive_jvps[scatter_max_p] = partial(_scatter_extremal_jvp, scatter_max_p)
 
 def _scatter_jvp(primals, tangents, *, update_jaxpr, update_consts,
@@ -2401,7 +2401,7 @@ def _scatter_transpose_rule(t, operand, indices, updates, *,
 ad.primitive_jvps[scatter_p] = _scatter_jvp
 ad.primitive_transposes[scatter_p] = _scatter_transpose_rule
 batching.primitive_batchers[scatter_p] = (
-  partial(_scatter_batching_rule, scatter))
+  partial(_scatter_batching_rule, scatter_p))
 
 
 def _scatter_lower_opaque(ctx, operand, indices, updates, *,


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -466,6 +466,15 @@ def jnp_op(x, idx):
     self._CheckAgainstNumpy(np_op_idx, jnp_op_idx, args_maker)
     self._CompileAndCheck(jnp_op_idx, args_maker)
 
+  def testIndexApplyBatchingBug(self):
+    # https://github.com/google/jax/issues/16655
+    arr = jnp.array([[1, 2, 3, 4, 5, 6]])
+    ind = jnp.array([3])
+    func = lambda a, i: a.at[i].apply(lambda x: x - 1)
+    expected = jnp.array(list(map(func, arr, ind)))
+    out = jax.vmap(func)(arr, ind)
+    self.assertArraysEqual(out, expected)
+
   def testIndexUpdateScalarBug(self):
     # https://github.com/google/jax/issues/14923
     a = jnp.arange(10.)
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2302,6 +2302,29 @@ def testScatterMax(self, arg_shape, dtype, idxs, update_shape, dnums):
     fun = partial(lax.scatter_max, dimension_numbers=dnums)
     self._CompileAndCheck(fun, args_maker)
 
+  @jtu.sample_product(
+    [dict(arg_shape=arg_shape, idxs=idxs, update_shape=update_shape, dnums=dnums)
+      for arg_shape, idxs, update_shape, dnums in [
+          ((5,), np.array([[0], [2]]), (2,), lax.ScatterDimensionNumbers(
+            update_window_dims=(), inserted_window_dims=(0,),
+            scatter_dims_to_operand_dims=(0,))),
+          ((10,), np.array([[0], [0], [0]]), (3, 2), lax.ScatterDimensionNumbers(
+            update_window_dims=(1,), inserted_window_dims=(),
+            scatter_dims_to_operand_dims=(0,))),
+          ((10, 5,), np.array([[0], [2], [1]]), (3, 3), lax.ScatterDimensionNumbers(
+            update_window_dims=(1,), inserted_window_dims=(0,),
+            scatter_dims_to_operand_dims=(0,))),
+    ]],
+    dtype=lax_test_util.float_dtypes,
+  )
+  def testScatterApply(self, arg_shape, dtype, idxs, update_shape, dnums):
+    rng = jtu.rand_default(self.rng())
+    rng_idx = jtu.rand_int(self.rng(), high=max(arg_shape))
+    rand_idxs = lambda: rng_idx(idxs.shape, idxs.dtype)
+    args_maker = lambda: [rng(arg_shape, dtype), rand_idxs()]
+    fun = partial(lax.scatter_apply, func=jnp.sin, update_shape=update_shape, dimension_numbers=dnums)
+    self._CompileAndCheck(fun, args_maker)
+
   @jtu.sample_product(
     [dict(arg_shape=arg_shape, idxs=idxs, update_shape=update_shape,
           dnums=dnums)
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -601,6 +601,29 @@ def testScatterAdd(self, arg_shape, dtype, idxs, update_shape, dnums, bdims):
                         [dtype, idxs.dtype, dtype], jtu.rand_default(self.rng()),
                         rtol={np.float16: 5e-3, dtypes.bfloat16: 7e-2})
 
+  @jtu.sample_product(
+    [dict(arg_shape=arg_shape, idxs=idxs, update_shape=update_shape,
+          dnums=dnums, bdims=bdims)
+      for arg_shape, idxs, update_shape, dnums in [
+          ((5,), np.array([[0], [2]]), (2,), lax.ScatterDimensionNumbers(
+            update_window_dims=(), inserted_window_dims=(0,),
+            scatter_dims_to_operand_dims=(0,))),
+          ((10,), np.array([[0], [0], [0]]), (3, 2), lax.ScatterDimensionNumbers(
+            update_window_dims=(1,), inserted_window_dims=(),
+            scatter_dims_to_operand_dims=(0,))),
+          ((10, 5,), np.array([[0], [2], [1]]), (3, 3), lax.ScatterDimensionNumbers(
+            update_window_dims=(1,), inserted_window_dims=(0,),
+            scatter_dims_to_operand_dims=(0,))),
+      ]
+      for bdims in lax_test_util.all_bdims(arg_shape, idxs.shape)],
+    dtype=lax_test_util.float_dtypes,
+  )
+  def testScatterApply(self, arg_shape, dtype, idxs, update_shape, dnums, bdims):
+    fun = partial(lax.scatter_apply, func=jnp.sin, update_shape=update_shape, dimension_numbers=dnums)
+    self._CheckBatching(fun, 5, bdims, [arg_shape, idxs.shape],
+                        [dtype, idxs.dtype], jtu.rand_default(self.rng()),
+                        rtol={np.float16: 5e-3, dtypes.bfloat16: 7e-2})
+
   def testShapeUsesBuiltinInt(self):
     x = lax.iota(np.int32, 3) + 1
     self.assertIsInstance(x.shape[0], int)  # not np.int64