Pooling layers

[gboduljak]

Proposed changes

Added MaxPool and AvgPool layers. MaxPool and AvgPool layers were requested in this issue. In this PR, I propose implementing the requested pooling operations by firstly computing sliding windows and subsequently reducing them. More precisely, we can use as_strided to compute pooling sliding windows. Then, we can simply reduce over appropriate axes to implement the desired pooling operation.

Concerns: My only concern is performance. Ideally, we should call optimized backend kernels for pooling operations. There are MPSCNNPoolingAverageNode and MPSCNNPoolingMaxNode. Similarly, there is BNNSFilterCreateLayerPooling in Accelerate. Alternatively, we could implement a kernel for window reduction.

Closes https://github.com/ml-explore/mlx/issues/25.

Checklist

Put an x in the boxes that apply.

• [x] I have read the CONTRIBUTING document
• [x] I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• [x] I have added tests that prove my fix is effective or that my feature works
• [x] I have updated the necessary documentation (if needed)

[gboduljak]

@awni @angeloskath Could you take a look at this draft implementation?

[gboduljak]

@angeloskath, do you have any performance concerns regarding pooling implementation using as_strided ?

[gboduljak]

@angeloskath I implemented the requested changes. In addition, I separated the generic implementation into MaxPooling1d, MaxPooling2d, AvgPooling1d and AvgPooling2d and updated the docs accordingly. However, I am not sure about my inheritance hierarchy.

I would also appreciate if you could take a detailed look into some of the tests. The expected results were obtained by computing outputs of torch.nn pooling layer implementations and transposing the outputs to match our channels_last convention. There could be some mistakes.

Please take another look :)

[gboduljak]

@angeloskath Could you take a look at the changes addressing your comments?

[RahulBhalley]

Hi guys! Any updates on this? Looks like Transformers get all the Attention (pun intended). 😅

Nearly no CNN can be implemented without pooling layers (even AlexNet, I'm trying to write VGG19 but pooling layers are needed). Can we prioritise this please???

P.S. I cannot use MLX at all until this gets merged. I am mostly into CNN.

CC: @angeloskath @awni @gboduljak

[gboduljak]

@RahulBhalley I would like to complete this PR :) I am waiting for the review.

[RahulBhalley]

I highly appreciate your efforts @gboduljak. I hope this gets reviewed ASAP.

[angeloskath]

@gboduljak I am coming back to this PR (sorry for taking too long). Could you perhaps put it back on top of main? Since rebase is pretty hard given the many merge commits I propose applying the diff on top of main and force pushing to the same branch. Something using the following

base=$(git merge-base main pooling-layers)
git checkout main
git checkout -b new-pooling-layers
git diff $base..pooling-layers >/tmp/pooling-layers.patch
git apply --reject /tmp/pooling-layers.patch
rm docs/src/python/nn/layers.rst.rej
git add <whatever is modified or uncommitted>
git commit -m 'Add pooling layers'
git branch -m pooling-layers old-pooling-layers
git branch -m pooling-layers
git push origin pooling-layers -f

I can do the above but then the authorship of the commit will be wrong (me instead of you).

[gboduljak]

@angeloskath Thank you for informing me. I will attempt the rebase today.

[gboduljak]

@angeloskath I did the rebase according to your instructions. Many thanks for the detailed instructions. Please take a look :)

[awni]

Looks great! A few issues in the docs, please fix and then I think we can merge.

[awni]

FYI some instructions on building the docs here. It's useful to build / look at them for big changes to the docs otherwise you can break them / make strange outputs.

[ligaz]

@awni It would be great if there is a preview environment where you can browse your changes to the docs. This will eliminate the need to this manually.

[awni]

It would be great if there is a preview environment where you can browse your changes to the docs. This will eliminate the need to this manually.

Interesting idea.. maybe something we can setup with CircleCI 🤔

[awni]

@gboduljak actually I just pushed the docs changes myself since I already made them locally.

LGTM, @angeloskath feel free to merge when you are good with it.

Thanks for the contribution!

[gboduljak]

@angeloskath Thank you for the refactor. @awni Thank you for fixing the docs nits. I was about to push my changes, but PR was already merged :)