flax

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