Repository: google/flax

Issue:
Pooling operations should support inputs with multiple batch dimensions.
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.
 
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04 (anaconda)
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: flax=0.6.1, jax=0.3.23, jaxlib=0.3.22+cuda11.cudnn82
- Python version: 3.10.6
- GPU/TPU model and memory: RTX3060, 12GB
- CUDA version (if applicable): 11.3


### Problem you have encountered:
In Flax, pooling operations (i.e., avg_pool and max_pool) only support arrays with shape (batch, window dims..., features) or (window dims..., features) as their inputs. If arrays with multiple batch dimensions are given, pooling operations raise AssertionError. 

However, many layers such as Dense, Conv, BatchNorm, etc. support inputs with multiple batch dimensions. For example, the docstring of Conv.__call__ explains that "If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning". I think supporting inputs with multiple batch dimensions is sometimes useful, and pooling operations should support such inputs like other layers. 

### What you expected to happen:
avg_pool and max_pool should support arrays with shape (*batch_dims, window dims..., features) where len(batch_dims)>2 as their inputs. 

### Logs, error messages, etc:
When I run the code in "steps to reproduce", the following error is raised.

Traceback (most recent call last):
  File "/home/hayato/Code/research/pool.py", line 15, in <module>
    y = linen.max_pool(x, (2, 2), (2, 2))
  File "/home/hayato/miniconda3/envs/jax/lib/python3.10/site-packages/flax/linen/pooling.py", line 113, in max_pool
    y = pool(inputs, -jnp.inf, lax.max, window_shape, strides, padding)
  File "/home/hayato/miniconda3/envs/jax/lib/python3.10/site-packages/flax/linen/pooling.py", line 57, in pool
    assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})"
AssertionError: len((1, 8, 32, 32, 3)) != len((1, 2, 2, 1))


### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.

```python
import jax.random as jr
import jax.numpy as jnp
from flax import linen

# Image-like array with multiple batch dimensions. batch_dims=(1, 8)
x = jnp.zeros((1, 8, 32, 32, 3), dtype=jnp.float32)

# Conv does not raise any errors.
y, _ = linen.Conv(16, (3, 3)).init_with_output(jr.PRNGKey(0), x) 
print(y.shape)  # (1, 8, 32, 32, 16)

# max_pool raises AssertionError.
y = linen.max_pool(x, (2, 2), (2, 2))  
print(y.shape)  # Expected output: (1, 8, 16, 16, 3)
```


Discussion:

Resolving patch:
diff --git a/flax/linen/pooling.py b/flax/linen/pooling.py
--- a/flax/linen/pooling.py
+++ b/flax/linen/pooling.py
@@ -41,17 +41,20 @@ def pool(inputs, init, reduce_fn, window_shape, strides, padding):
   Returns:
     The output of the reduction for each window slice.
   """
+  num_batch_dims = inputs.ndim - (len(window_shape) + 1)
   strides = strides or (1,) * len(window_shape)
   assert len(window_shape) == len(strides), (
       f"len({window_shape}) must equal len({strides})")
-  strides = (1,) + strides + (1,)
-  dims = (1,) + window_shape + (1,)
+  strides = (1,) * num_batch_dims + strides + (1,)
+  dims = (1,) * num_batch_dims + window_shape + (1,)
 
   is_single_input = False
-  if inputs.ndim == len(dims) - 1:
+  if num_batch_dims == 0:
     # add singleton batch dimension because lax.reduce_window always
     # needs a batch dimension.
     inputs = inputs[None]
+    strides = (1,) + strides
+    dims = (1,) + dims
     is_single_input = True
 
   assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})"


Test patch:
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -119,6 +119,17 @@ def test_avg_pool_padding_same(self, count_include_pad):
       expected_y = jnp.array([10.0 / 4, 6.0 / 2, 7.0 / 2, 4.0 / 1]).reshape((1, 2, 2, 1))
     np.testing.assert_allclose(y, expected_y)
 
+  def test_pooling_variable_batch_dims(self):
+    x = jnp.zeros((1, 8, 32, 32, 3), dtype=jnp.float32)
+    y = nn.max_pool(x, (2, 2), (2, 2))
+
+    assert y.shape == (1, 8, 16, 16, 3)
+
+  def test_pooling_no_batch_dims(self):
+    x = jnp.zeros((32, 32, 3), dtype=jnp.float32)
+    y = nn.max_pool(x, (2, 2), (2, 2))
+
+    assert y.shape == (16, 16, 3)
 
 class NormalizationTest(parameterized.TestCase):