Repository: google/jax

Issue:
`jax.vmap(jax.pure_callback(...), in_axes=1)` is broken
### Description

When `vmap`ping a `jax.pure_callback` over the non-0 axis, the logic that checks the callback output shape seems wrong.

MWE:
```python
from functools import partial

import jax
import jax.numpy as jnp
import numpy as np

def my_fun(x):
    print(f"provided callback with {x.shape}")
    return x

@jax.jit
def my_jax_fun(x):
    res = jax.pure_callback(my_fun, 
                      jax.ShapeDtypeStruct(x.shape, x.dtype),
                      x,
                      vectorized=True)
    return res

s = jnp.ones((2,3,4))
vmap_fun = jax.vmap(my_jax_fun, in_axes=1, out_axes=1)
vmap_fun(s)
```

error:
```python
provided callback with (3, 2, 4)
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Cell In[13], line 23
     19 s = jnp.ones((2,3,4))
     21 vmap_fun = jax.vmap(my_jax_fun, in_axes=1, out_axes=1)
---> 23 vmap_fun(s)

    [... skipping hidden 17 frame]

File ~/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py:1201, in ExecuteReplicated.__call__(self, *args)
   1198 if (self.ordered_effects or self.has_unordered_effects
   1199     or self.has_host_callbacks):
   1200   input_bufs = self._add_tokens_to_inputs(input_bufs)
-> 1201   results = self.xla_executable.execute_sharded(
   1202       input_bufs, with_tokens=True
   1203   )
   1204   result_token_bufs = results.disassemble_prefix_into_single_device_arrays(
   1205       len(self.ordered_effects))
   1206   sharded_runtime_token = results.consume_token()

XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: Incorrect output shape for return value 0: Expected: (2, 2, 4), Actual: (3, 2, 4)
```

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.24
jaxlib: 0.4.24
numpy:  1.24.3
python: 3.11.2 (main, Apr  7 2023, 16:35:55) [Clang 14.0.3 (clang-1403.0.22.14.1)]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
```

cc @Adrien-Kahn


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
@@ -97,7 +97,7 @@ def pure_callback_batching_rule(
     vectorized: bool,
     result_avals: Sequence[core.ShapedArray],
 ):
-  axis_size = next(a.shape[0] for a, d in zip(args, dims)
+  axis_size = next(a.shape[d] for a, d in zip(args, dims)
                    if d is not batching.not_mapped)
   new_args = [arg if dim is batching.not_mapped else
               batching.moveaxis(arg, dim, 0) for arg, dim in zip(args, dims)]


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
@@ -566,6 +566,16 @@ def h(x, y):
     self.assertArraysAllClose(out, np.sin(np.arange(4.)) + np.arange(10., 14.),
                               rtol=1E-7, check_dtypes=False)
 
+    @jax.jit
+    @functools.partial(jax.vmap, in_axes=1, out_axes=1)
+    def h(x, y):
+      out_shape = jax.ShapeDtypeStruct(x.shape, np.result_type(x.dtype, y.dtype))
+      return jax.pure_callback(lambda x, y: np.sin(x) + y, out_shape, x, y)
+    out = h(jnp.arange(4.)[None], jnp.arange(10., 14.)[None])
+    self.assertArraysAllClose(out, np.sin(np.arange(4.)) + np.arange(10.,
+      14.)[None],
+                              rtol=1E-7, check_dtypes=False)
+
   def test_vmap_vectorized_callback(self):
 
     def cb(x):
@@ -598,6 +608,15 @@ 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, in_axes=(1, None), out_axes=1)
+    def h(x, y):
+      return jax.pure_callback(lambda x, y: np.sin(x) + y, x, x, y,
+                               vectorized=True)
+    out = h(jnp.arange(4.)[None], 4.)
+    np.testing.assert_allclose(out, np.sin(np.arange(4.)[None]) + 4.)
+
+
   def test_vmap_vectorized_callback_errors_if_returns_wrong_shape(self):
 
     def cb(x):