Repository: google/jax

Issue:
Vmap of pure_callback with multiple arguments error
### Description

Here is the minimal example of the error:
```
import jax.numpy as jnp
from jax import pure_callback, vmap

dim = 5
num_samples = 3
test_a = jnp.ones(shape=(num_samples, dim))


def _simple_pass(a):
    return a


def simple_pass(a):
    return pure_callback(_simple_pass, a, a, vectorized=False)


out = vmap(simple_pass)(test_a)


def _simple_pass_pair(a, b):
    return a, b


def simple_pass_pair(a, b):
    return pure_callback(_simple_pass_pair, (a, b), a, b, vectorized=False)


out_pair = vmap(simple_pass_pair)(test_a, test_a)
```
While the vmapping in the case when we have one input argument works, in the case when we have **two or more** input arguments we get the following error:

```
Traceback (most recent call last):
  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/code.py", line 90, in runcode
    exec(code, self.locals)
  File "<input>", line 1, in <module>
  File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 198, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/lenarttreven/PycharmProjects/learning_dynamical_systems/experiments/random/pure_callback.py", line 28, in <module>
    out_pair = vmap(simple_pass_pair)(test_a, test_a)
  File "/Users/lenarttreven/PycharmProjects/learning_dynamical_systems/experiments/random/pure_callback.py", line 25, in simple_pass_pair
    return pure_callback(_simple_pass, (a, b), a, b, vectorized=False)
  File "/Users/lenarttreven/python_venv/artemis/lib/python3.10/site-packages/jax/_src/callback.py", line 136, in pure_callback
    out_flat = pure_callback_p.bind(
  File "/Users/lenarttreven/python_venv/artemis/lib/python3.10/site-packages/jax/_src/callback.py", line 95, in pure_callback_batching_rule
    outvals = lax_map(_batch_fun, *batched_args)
TypeError: map() takes 2 positional arguments but 3 were given
``` 


### What jax/jaxlib version are you using?

jax v0.3.20, jaxlib v0.3.20

### Which accelerator(s) are you using?

CPU

### Additional system info

Mac M1

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -86,13 +86,13 @@ def pure_callback_batching_rule(args, dims, *, callback, vectorized: bool,
   else:
     is_batched = [d is not batching.not_mapped for d in dims]
     unbatched_args, batched_args = util.partition_list(is_batched, new_args)
-    def _batch_fun(*batched_args):
+    def _batch_fun(batched_args):
       merged_args = util.merge_lists(is_batched, unbatched_args, batched_args)
       return pure_callback_p.bind(
           *merged_args, callback=callback, result_avals=result_avals,
           vectorized=vectorized)
     from jax._src.lax.control_flow import map as lax_map
-    outvals = lax_map(_batch_fun, *batched_args)
+    outvals = lax_map(_batch_fun, batched_args)
   return tuple(outvals), (0,) * len(outvals)
 
 


Test patch:
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -592,6 +592,13 @@ def h(x, y):
     out = h(jnp.arange(4.), 4.)
     np.testing.assert_allclose(out, np.sin(np.arange(4.)) + 4.)
 
+    @jax.jit
+    @functools.partial(jax.vmap)
+    def h(x, y):
+      return jax.pure_callback(lambda x, y: np.sin(x) + y, x, x, y)
+    out = h(jnp.arange(4.), jnp.arange(10., 14.))
+    np.testing.assert_allclose(out, np.sin(np.arange(4.)) + jnp.arange(10., 14.))
+
   def test_vmap_vectorized_callback(self):
 
     def cb(x):