flax's avg_pool doesn't play well with vmap

[bilal2vec]

Hi!

I'm using flax's avg pooling and i've run into an error when trying to use it inside a vmaped function:

@jax.vmap
def apply(x):
  return nn.avg_pool(x, window_shape=(3, 3), strides=(2, 2), padding=((1, 1), (1, 1)))

TypeError: reduce_window got the wrong number of window_dimensions for operand: got operand shape (8, 32, 32, 3) with window_dimensions (1, 1, 3, 3, 1).

idk if it isn't supposed to be used in this way because the imagenet and mnist examples use it without wrapping everything in a vmap.

I have a colab notebook with the full stacktrace as well

[jheek]

Yes I think you are right. Pooling layers still assume there is a batch dimension. This should be an easy fix. I created an issue for this #739