Regression bug in associative_scan
Hi,

There has been a regression in JAX code, which I believe to be related to this [commit](https://github.com/google/jax/commit/d3db7bd4be96e1fc16daa4571ade7fbfce13197f#diff-cc3f3a1265e2fee9115dc5a9f619d402e280fb48a7ed7d91261b42813d200f1b). 

I have created PR https://github.com/google/jax/pull/5165 with a reproducing failing test.

I will try and look further into what the problem could be, but if the person who made said commit could also have a look it would be great (@hawkinsp I believe this is you).

Adrien
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 557a47acb..6183084f2 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4352,7 +4352,7 @@ def _gather_batching_rule(batched_args, batch_dims, *, dimension_numbers,
     counts = broadcasted_iota(start_indices.dtype, tuple(count_shape), 0)
     start_indices = concatenate([counts, start_indices], len(count_shape) - 1)
 
-    slice_sizes = (1,) + slice_sizes
+    slice_sizes = (_min(operand.shape[0], 1),) + slice_sizes
     collapsed_slice_dims = (0,) + tuple(np.add(1, dimension_numbers.collapsed_slice_dims))
     offset_dims = tuple(np.add(1, dimension_numbers.offset_dims))
     start_index_map = (0,) + tuple(np.add(1, dimension_numbers.start_index_map))
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
index d38d4c9df..2c847eecf 100644
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -74,7 +74,12 @@ class LaxVmapTest(jtu.JaxTestCase):
     args = [rng(shape, dtype) for shape, dtype in zip(batched_shapes, dtypes)]
     args_slice = args_slicer(args, bdims)
     ans = api.vmap(op, bdims)(*args)
-    expected = np.stack([op(*args_slice(i)) for i in range(bdim_size)])
+    if bdim_size == 0:
+      args = [rng(shape, dtype) for shape, dtype in zip(shapes, dtypes)]
+      out = op(*args)
+      expected = np.zeros((0,) + out.shape, out.dtype)
+    else:
+      expected = np.stack([op(*args_slice(i)) for i in range(bdim_size)])
     self.assertAllClose(ans, expected, rtol=rtol, atol=atol)
 
   @parameterized.named_parameters(itertools.chain.from_iterable(
@@ -642,6 +647,8 @@ class LaxVmapTest(jtu.JaxTestCase):
       for bdims in all_bdims(shape, idxs.shape)))
   def testGather(self, shape, dtype, idxs, dnums, slice_sizes, bdims):
     fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    self._CheckBatching(fun, 0, bdims, [shape, idxs.shape], [dtype, idxs.dtype],
+                        jtu.rand_default(self.rng()))
     self._CheckBatching(fun, 5, bdims, [shape, idxs.shape], [dtype, idxs.dtype],
                         jtu.rand_default(self.rng()))
 
['tests/lax_vmap_test.py', 'jax/_src/lax/lax.py']