Repository: google/jax

Issue:
`a.at[i].apply(...)` fails when `i: int` 
### Description

```python
a = jnp.arange(10.)
a = a.at[0].apply(jnp.sin)
```
fails with
```python
TypeError: Updates tensor must be of rank 0; got 1.
```

I think it should work, as `at[0:1]` works. The error message is also not that helpful. It can produce more cryptic error (`MLIRError`) when `a.dtype == np.int_`. NumPy seems to cast the result back to int:

```python
a = np.arange(10)
np.cos.at(a, np.arange(3))
# a == array([1, 0, 0, 3, 4, 5, 6, 7, 8, 9])
```
Not sure we want the same behaviour, but I think we could at least raise a better error than `MLIRError`.


Discussion:

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
@@ -514,6 +514,7 @@ def scatter_apply(
   operand: Array, scatter_indices: Array,
   func: Callable[[Array], Array],
   dimension_numbers: ScatterDimensionNumbers, *,
+  update_shape: Shape = (),
   indices_are_sorted: bool = False, unique_indices: bool = False,
   mode: Optional[Union[str, GatherScatterMode]] = None) -> Array:
   """Scatter-apply operator.
@@ -539,6 +540,7 @@ def scatter_apply(
     dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
       how dimensions of `operand`, `start_indices`, `updates` and the output
       relate.
+    update_shape: the shape of the updates at the given indices.
     indices_are_sorted: whether `scatter_indices` is known to be sorted. If
       true, may improve performance on some backends.
     unique_indices: whether the elements to be updated in ``operand`` are
@@ -555,7 +557,7 @@ def scatter_apply(
     An array containing the result of applying `func` to `operand` at the given indices.
   """
   # TODO: can we implement this without a placeholder?
-  unused = lax.full(scatter_indices.shape[:1], 0, operand.dtype)
+  unused = lax.full(update_shape, 0, operand.dtype)
   _apply = lambda x, _: func(x)
   try:
     _apply = _scatter_apply_cache.setdefault(func, _apply)
diff --git a/jax/_src/numpy/array_methods.py b/jax/_src/numpy/array_methods.py
--- a/jax/_src/numpy/array_methods.py
+++ b/jax/_src/numpy/array_methods.py
@@ -541,8 +541,8 @@ def apply(self, func, *, indices_are_sorted=False, unique_indices=False,
 
     See :mod:`jax.ops` for details.
     """
-    def _scatter_apply(x, indices, _, dims, **kwargs):
-      return lax.scatter_apply(x, indices, func, dims, **kwargs)
+    def _scatter_apply(x, indices, y, dims, **kwargs):
+      return lax.scatter_apply(x, indices, func, dims, update_shape=y.shape, **kwargs)
     return scatter._scatter_update(self.array, self.index,
                                    lax_internal._zero(self.array.dtype),
                                    _scatter_apply,


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
@@ -452,10 +452,25 @@ def np_op(x, idx):
       return y
     def jnp_op(x, idx):
       return jnp.asarray(x).at[idx].apply(jnp_func)
+
+    # Test with traced integer index
     args_maker = lambda: [rng(size, dtype), idx_rng(size, int)]
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+    # Test with slice index
+    idx = slice(1, 5)
+    np_op_idx = partial(np_op, idx=idx)
+    jnp_op_idx = partial(jnp_op, idx=idx)
+    args_maker = lambda: [rng(size, dtype)]
+    self._CheckAgainstNumpy(np_op_idx, jnp_op_idx, args_maker)
+    self._CompileAndCheck(jnp_op_idx, args_maker)
+
+  def testIndexUpdateScalarBug(self):
+    # https://github.com/google/jax/issues/14923
+    a = jnp.arange(10.)
+    out = a.at[0].apply(jnp.cos)
+    self.assertArraysEqual(out, a.at[0].set(1))
 
   @jtu.sample_product(
     [dict(name=name, shape=shape, indexer=indexer, mode=mode)