Implementing PyTorch's AdaptiveAvgPooling2D in Flax

[muhd-umer]

Hi, so I have been trying to implement PyTorch's AdaptiveAvgPooling2d in Flax and from their official docs and internet sources, I came to these formulas to compute the parameter values to adapt to the input.

stride = (input_size // output_size) 
kernel_size = input_size - (output_size - 1)*stride  
padding = 0

Now while this may not be the right place to ask and validate my implementation, I am not really sure if I can take advantage of nn.avg_pool to implement the adaptive pooling layer. Following is what I came up with;

class AdaptiveAveragePool2D(nn.Module):
    output_size: Union[int, Iterable[int]]

    @nn.compact
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        assert (
            x.ndim > 2 and x.ndim < 5
        ), f"Expected input to have 3 (or 4 with batch) but recieved {x.ndim} dims."
        if x.ndim == 4:
            input_size = x.shape[1], x.shape[2]
        else:
            input_size = x.shape[0], x.shape[1]
        if not isinstance(self.output_size, Iterable):
            output_size = to_2tuple(self.output_size)
        else:
            output_size = self.output_size

        strides = tuple(e1 // e2 for e1, e2 in zip(input_size, output_size))
        factor = tuple((e1 - 1) * e2 for e1, e2 in zip(output_size, strides))
        window_shape = tuple(e1 - e2 for e1, e2 in zip(input_size, factor))

        x = nn.avg_pool(x, window_shape=window_shape, strides=strides, padding="VALID")

        return x

While the desired output shape is obtained correctly, I'm unsure if there are any cases for which the above might fail.
Any help is appreciated ❤️.