Pooling operation

[jorisSchaller]

Moved pooling function out of linen according to the comment #4408 (comment)

What does this PR do?

It moves the pooling function from linen to flax

I moved the pooling operation as they can be better shared between Linen and nnx.
The documentation was also missing from the nnx documentation. This PR adds an element in the toc as well as a page describing the function

Fixes #4271
closes #4408

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).

[vfdev-5]

@jorisSchaller thanks for the PR, can you please rebase it on the latest main to avoid unrelated changes? Thanks!

[jorisSchaller]

@vfdev-5 let me know if you'd like any additional changes

[jorisSchaller]

@vfdev-5 reverted changes and merged upstream changes

[vfdev-5]

@jorisSchaller please check this failing test: https://github.com/google/flax/actions/runs/18982441063/job/54231421369?pr=5057

[jorisSchaller]

@vfdev-5 Updated

[vfdev-5]

@jorisSchaller thanks for updates on the PR !
We still need to fix few things:

• squash the commit history to 1 commit. Max 5 commits per PR.
• fix this failing pre-commit hooks test: https://github.com/google/flax/actions/runs/19017841913/job/54308506195?pr=5057

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[vfdev-5]

@jorisSchaller please rebase on main and fix the failing check with ruff: https://github.com/google/flax/actions/runs/19210434754/job/54950332158?pr=5057

[vfdev-5]

Thanks @jorisSchaller for PR. I finished the remaining things and now it LGTM