Repository: google/jax
Issue: AvgPool does only work as global average pooling

```python3
from jax.experimental import stax
import jax.numpy as np

init_fn, apply_fun = stax.AvgPool((2, 2), strides=(2, 2))
output_shape, params = init_fn((-1, 32, 32, 3))
print(output_shape)
print(params)
apply_fun(params, np.zeros((100, 32, 32, 3)))
```

This should work but instead it fails with this error:
```
(-1, 16, 16, 3)
()
Traceback (most recent call last):
  File "minimal_example.py", line 9, in <module>
    apply_fun(params, np.zeros((100, 32, 32, 3)))
  File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/experimental/stax.py", line 172, in apply_fun
    return rescale(out, inputs) if rescale else out
  File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/experimental/stax.py", line 183, in rescale
    return outputs / window_sizes
  File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 350, in true_divide
    x1, x2 = _promote_shapes(x1, x2)
  File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 134, in _promote_shapes
    nd = len(_broadcast_shapes(*shapes))
  File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/util.py", line 161, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 151, in _broadcast_shapes
    .format(tuple(map(tuple, shapes))))
ValueError: Incompatible shapes for broadcasting: ((100, 16, 16, 3), (1, 1, 16, 16))
```

Note that it also doesn't work if stride is 1. The only case in which `AvgPool` works seems to be if stride is 1 and the pooling size is identical to the input size (i.e. global average pooling). `MaxPool` seems to work fine.

Changed files:
jax/experimental/stax.py, tests/stax_test.py

diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
index 4a2e4abb7..373f9f1b7 100644
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -178,9 +178,9 @@ SumPool = _pooling_layer(lax.add, 0.)
 
 def _normalize_by_window_size(dims, strides, padding):
   def rescale(outputs, inputs):
-    one = np.ones(inputs.shape[1:3], dtype=inputs.dtype)
+    one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
     window_sizes = lax.reduce_window(one, 0., lax.add, dims, strides, padding)
-    return outputs / window_sizes
+    return outputs / window_sizes[..., np.newaxis]
   return rescale
 AvgPool = _pooling_layer(lax.add, 0., _normalize_by_window_size)
 
diff --git a/tests/stax_test.py b/tests/stax_test.py
index 245c9b24e..e55e06dc3 100644
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -100,16 +100,20 @@ class StaxTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_window_shape={}_padding={}_strides={}_input_shape={}"
-                        .format(window_shape, padding, strides, input_shape),
+                        "_maxpool={}"
+                        .format(window_shape, padding, strides, input_shape,
+                                max_pool),
        "window_shape": window_shape, "padding": padding, "strides": strides,
-       "input_shape": input_shape}
+       "input_shape": input_shape, "max_pool": max_pool}
       for window_shape in [(1, 1), (2, 3)]
       for padding in ["VALID"]
       for strides in [None, (2, 1)]
-      for input_shape in [(2, 5, 6, 1)]))
-  def testPoolingShape(self, window_shape, padding, strides, input_shape):
-    init_fun, apply_fun = stax.MaxPool(window_shape, padding=padding,
-                                       strides=strides)
+      for input_shape in [(2, 5, 6, 1)]
+      for max_pool in [False, True]))
+  def testPoolingShape(self, window_shape, padding, strides, input_shape,
+                       max_pool):
+    layer = stax.MaxPool if max_pool else stax.AvgPool
+    init_fun, apply_fun = layer(window_shape, padding=padding, strides=strides)
     _CheckShapeAgreement(self, init_fun, apply_fun, input_shape)
 
   @parameterized.named_parameters(jtu.cases_from_list(