How to Implement MaxPooling2D with Ceil Mode in Flax?

[firqaaa]

What is the equivalent code for maxpooling2d with ceil mode in PyTorch?

So for example in PyTorch we can code like this :
• nn.MaxPool2d(3, stride=2, padding=0, ceil_mode=True)
But in Flax, the framework only supports four parameters such as input, window_size, strides, and padding.

[andsteing]

Note that Flax also exports the more powerful nn.pool() function; this can be used for example to implement torch.nn.AvgPool2d as described here:
#1916

[marcvanzee]

Hi @firqaaa, we currently don't support ceil mode for our pooling layers. The best way to reproduce pytorch behavior is to use VALID padding and then pad the input manually to the shapes you want using the padding argument.