Enhancement: 2d Pooling

[torridgristle]

Could put in 2d pooling for an easy to drop in alternative to AvgPool2d by using pixel_unshuffle for cases where there's exact divisions of an input shape, or padding and unfold for overlapping windows / strides that aren't equal to the window size.

I don't know how fast Unfold is or if there's a better way to do the operation; I tried checking the PyTorch github to see how it does AvgPool2d for strides that aren't equal to the window size behind the scenes but I can never seem to figure out which version of the operation is the one that's used, it's defined in so many places it's beyond me.

I've once seen an alternative to pixel unshuffle that used grouped conv2d, a kernel that put each position into its own channel output, and stride, but I can't seem to find it again. It was useful because you could adjust the stride and then it'd be like Unfold, but I never got around to testing if it was actually faster.

[lucidrains]

@torridgristle sure i'll think about it, perhaps can build it with https://github.com/openai/triton

[lucidrains]

@torridgristle would you be willing to share some of your experimental results, if i were to make the LogAvgExpPool2d?

[lucidrains]

@torridgristle https://github.com/lucidrains/logavgexp-torch/tree/0.0.3#logavgexp2d

[ekurtulus]

Thanks a lot @lucidrains, a great implementation. Will 3D variant be implemented as well ?

[lucidrains]

@ekurtulus 👋 sure, why not 250852d

[ekurtulus]

Thanks a lot : )