Repository: google/jax

Issue:
 Issue with JAX's frompyfunc and at methods when compared to NumPy
### Description

I encountered an issue when trying to use JAX's frompyfunc and at methods for a specific use case that works fine in NumPy. Below is the code snippet that demonstrates the issue: 
```
import jax.numpy as jnp

def scalar_add(x, y):
  # emphasize that only scalar tracers will be passed to this function.
  assert jnp.shape(x) == jnp.shape(y) == ()
  return x + y

add = jnp.frompyfunc(scalar_add, nin=2, nout=1, identity=0)

x = jnp.ones((5, 3))
indices = jnp.array([0,4,2])
t = jnp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

x = add.at(x, indices, t, inplace=False) 
```
Expected Behavior:

The code above is expected to perform an addition operation using JAX's frompyfunc and at methods, similar to the behavior in NumPy.

Actual Behavior:

However, running this code results in the following error:
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[<ipython-input-139-8af5f01fcded>](https://localhost:8080/#) in <cell line: 19>()
     17 
     18 # add.accumulate(x)  # accumulate() method is cumulative reduction
---> 19 z = add.at(x, indices, t, inplace=False)  # at() method is similar to JAX's ndarray.at

    [... skipping hidden 12 frame]

3 frames
[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/ufunc_api.py](https://localhost:8080/#) in at(self, a, indices, b, inplace)
    238       return self._at_via_scan(a, indices)
    239     else:
--> 240       return self._at_via_scan(a, indices, b)
    241 
    242   def _at_via_scan(self, a, indices, *args):

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/ufunc_api.py](https://localhost:8080/#) in _at_via_scan(self, a, indices, *args)
    254       return a.at[indices].set(self._call(a.at[indices].get(), *args))
    255 
--> 256     args = tuple(_broadcast_to(arg, shape).ravel() for arg in args)
    257     indices = [idx if isinstance(idx, slice) else _broadcast_to(idx, shape).ravel() for idx in indices]
    258 

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/ufunc_api.py](https://localhost:8080/#) in <genexpr>(.0)
    254       return a.at[indices].set(self._call(a.at[indices].get(), *args))
    255 
--> 256     args = tuple(_broadcast_to(arg, shape).ravel() for arg in args)
    257     indices = [idx if isinstance(idx, slice) else _broadcast_to(idx, shape).ravel() for idx in indices]
    258 

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/util.py](https://localhost:8080/#) in _broadcast_to(arr, shape)
    397     shape_tail = shape[nlead:]
    398     compatible = all(core.definitely_equal_one_of_dim(arr_d, [1, shape_d])
--> 399                      for arr_d, shape_d in safe_zip(arr_shape, shape_tail))
    400     if nlead < 0 or not compatible:
    401       msg = "Incompatible shapes for broadcasting: {} and requested shape {}"

ValueError: safe_zip() argument 2 is shorter than argument 1
```

Additional Information:

This code works as expected in NumPy using the equivalent NumPy functions. The issue appears to be specific to JAX's implementation.  Numpy seems to work fine with the following results.
```
array([[ 2.,  3.,  4.],
       [ 1.,  1.,  1.],
       [ 8.,  9., 10.],
       [ 1.,  1.,  1.],
       [ 5.,  6.,  7.]])
```
Environment:

JAX version: 0.4.16
Python version: 0.4.16+cuda11.cudnn86
Operating System: Colab
Please let me know if any further information or clarification is needed to address this issue.

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

jax: 0.4.16, jaxlib: 0.4.16+cuda11.cudnn86

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

GPU T4 

### Additional system info

On Colab

### NVIDIA GPU info

nvidia-smi


Discussion:
Thanks for the report! I'm looking into it

Resolving patch:
diff --git a/jax/_src/numpy/ufunc_api.py b/jax/_src/numpy/ufunc_api.py
--- a/jax/_src/numpy/ufunc_api.py
+++ b/jax/_src/numpy/ufunc_api.py
@@ -20,6 +20,7 @@
 >>> arr = jnp.add.at(arr, ind, val, inplace=False)
 """
 from functools import partial
+import math
 import operator
 from typing import Any, Callable, Optional
 
@@ -264,6 +265,7 @@ def at(self, a: ArrayLike, indices: Any, b: Optional[ArrayLike] = None, /, *,
       return self._at_via_scan(a, indices, b)
 
   def _at_via_scan(self, a: ArrayLike, indices: Any, *args: Any) -> Array:
+    assert len(args) in {0, 1}
     check_arraylike(f"{self.__name__}.at", a, *args)
     dtype = jax.eval_shape(self._func, lax_internal._one(a), *(lax_internal._one(arg) for arg in args)).dtype
     a = lax_internal.asarray(a).astype(dtype)
@@ -277,7 +279,9 @@ def _at_via_scan(self, a: ArrayLike, indices: Any, *args: Any) -> Array:
     if not shape:
       return a.at[indices].set(self._call(a.at[indices].get(), *args))
 
-    args = tuple(_broadcast_to(arg, shape).ravel() for arg in args)
+    if args:
+      arg = _broadcast_to(args[0], (*shape, *args[0].shape[len(shape):]))
+      args = (arg.reshape(math.prod(shape), *args[0].shape[len(shape):]),)
     indices = [idx if isinstance(idx, slice) else _broadcast_to(idx, shape).ravel() for idx in indices]
 
     def scan_fun(carry, x):


Test patch:
diff --git a/tests/lax_numpy_ufuncs_test.py b/tests/lax_numpy_ufuncs_test.py
--- a/tests/lax_numpy_ufuncs_test.py
+++ b/tests/lax_numpy_ufuncs_test.py
@@ -289,6 +289,19 @@ def np_fun(x, idx, y):
     self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  def test_at_broadcasting(self):
+    # Regression test for https://github.com/google/jax/issues/18004
+    args_maker = lambda: [np.ones((5, 3)), np.array([0, 4, 2]),
+                          np.arange(9.0).reshape(3, 3)]
+    def np_fun(x, idx, y):
+      x_copy = np.copy(x)
+      np.add.at(x_copy, idx, y)
+      return x_copy
+    jnp_fun = partial(jnp.frompyfunc(jnp.add, nin=2, nout=1, identity=0).at, inplace=False)
+
+    self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
+    self._CompileAndCheck(jnp_fun, args_maker)
+
   @jtu.sample_product(
       SCALAR_FUNCS,
       [{'shape': shape, 'axis': axis}