Add count_include_pad argument to flax.linen.pooling.avg_pool / Fix div_shape for it raises error when there's no batch dimension
Add count_include_pad argument to flax.linen.pooling.avg_pool / Fix div_shape for it raises error when there's no batch dimension

google/flax,google/flax