Repository: google/jax

Issue:
Cannot vmap variadic lax.reduce_window  
### Discussed in https://github.com/google/jax/discussions/9818
@mattjj @hawkinsp @YouJiacheng

Minimal code to reproduce
```python
import jax
import jax.numpy as jnp
from jax import lax


def index_maxpool(x):
    def g(a, b):
        an, ai = a
        bn, bi = b
        which = an >= bn
        return (jnp.where(which, an, bn), jnp.where(which, ai, bi))

    _, idxs = lax.reduce_window(
        (x, jnp.arange(x.shape[0])),
        (-jnp.inf, -1),
        g,
        window_dimensions=(2,),
        window_strides=(2,),
        padding=((0, 0),),
    )
    return idxs

index_maxpool(jnp.ones((10,)))
jax.vmap(index_maxpool)(jnp.ones((10, 10)))
```

Error message
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_26237/2064479805.py in <module>
     22 
     23 index_maxpool(jnp.ones((10,)))
---> 24 jax.vmap(index_maxpool)(jnp.ones((10, 10)))

    [... skipping hidden 3 frame]

/tmp/ipykernel_26237/2064479805.py in index_maxpool(x)
     11         return (jnp.where(which, an, bn), jnp.where(which, ai, bi))
     12 
---> 13     _, idxs = lax.reduce_window(
     14         (x, jnp.arange(x.shape[0])),
     15         (-jnp.inf, -1),

    [... skipping hidden 4 frame]

~/base/lib/python3.9/site-packages/jax/_src/lax/windowed_reductions.py in _generic_reduce_window_batch_rule(batched_args, batch_dims, jaxpr, consts, window_dimensions, window_strides, padding, base_dilation, window_dilation)
    279   operand_bdims, init_value_bdims = util.split_list(batch_dims, [num_operands])
    280 
--> 281   operand, init = batched_args
    282   bdim, init_bdim = batch_dims
    283   if any(init_bdim is not None for init_bdim in init_value_bdims):

ValueError: too many values to unpack (expected 2)
```

Note: there is as well an issue when running that code on GPU (see @YouJiacheng comment in the Discussion thread)


Discussion:
This might be an easy one! I think Peter basically already solved this in his original commit, except [these two lines were accidentally left in](https://github.com/google/jax/commit/5415306257faede98f6c45f0ce67aeb1cd26b8ff#diff-9abc3334d4a9d008767f88a5df1b9d3fa2b176d5c168c6e47a7d3bcd2c844480R279-R280). That plus not adding tests for this case :P

Resolving patch:
diff --git a/jax/_src/lax/windowed_reductions.py b/jax/_src/lax/windowed_reductions.py
--- a/jax/_src/lax/windowed_reductions.py
+++ b/jax/_src/lax/windowed_reductions.py
@@ -278,8 +278,6 @@ def _generic_reduce_window_batch_rule(
   operands, init_values = util.split_list(batched_args, [num_operands])
   operand_bdims, init_value_bdims = util.split_list(batch_dims, [num_operands])
 
-  operand, init = batched_args
-  bdim, init_bdim = batch_dims
   if any(init_bdim is not None for init_bdim in init_value_bdims):
     raise NotImplementedError("reduce_window batching is not implemented for "
                               "initial values")


Test patch:
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -24,6 +24,7 @@
 import numpy as np
 
 import jax
+import jax.numpy as jnp
 from jax import dtypes
 from jax import lax
 
@@ -790,6 +791,37 @@ def testSort(self, shape, dimension, arity, bdims, is_stable):
   # TODO Collapse
   # TODO Scatter
 
+  # TODO(b/183233858): variadic reduce-window is not implemented on XLA:GPU
+  @jtu.skip_on_devices("gpu")
+  def test_variadic_reduce_window(self):
+    # https://github.com/google/jax/discussions/9818 and
+    # https://github.com/google/jax/issues/9837
+    def normpool(x):
+      norms = jnp.linalg.norm(x, axis=-1)
+      idxs = jnp.arange(x.shape[0])
+
+      def g(a, b):
+        an, ai = a
+        bn, bi = b
+        which = an >= bn
+        return (jnp.where(which, an, bn), jnp.where(which, ai, bi))
+
+      _, idxs = lax.reduce_window((norms, idxs), (-np.inf, -1), g,
+                        window_dimensions=(2,), window_strides=(2,),
+                        padding=((0, 0),))
+      return x[idxs]
+
+
+    inpt = jnp.array([
+      [1.0, 0.0, 1.0],
+      [2.0, 2.0, 0.0],
+      [3.0, 0.0, 1.0],
+      [0.0, 1.0, 1.0],
+    ])
+    output = jax.vmap(normpool)(inpt[None, ...])  # doesn't crash
+    expected = jnp.array([[[2.0, 2.0, 0.0], [3.0, 0.0, 1.0]]])
+    self.assertAllClose(output, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())