Merge pull request #19986 from sharadmv:pure-callback-bug

PiperOrigin-RevId: 610563393

Author: jax authors

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)]

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):