Skip to content

Pooling operation #5057

Merged
copybara-service[bot] merged 1 commit into
google:mainfrom
jorisSchaller:feat/nnx_pooling
Dec 9, 2025
Merged

Pooling operation #5057
copybara-service[bot] merged 1 commit into
google:mainfrom
jorisSchaller:feat/nnx_pooling

Conversation

@jorisSchaller

@jorisSchaller jorisSchaller commented Oct 28, 2025

Copy link
Copy Markdown
Contributor

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

vfdev-5 commented Oct 29, 2025

Copy link
Copy Markdown
Collaborator

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

@jorisSchaller

jorisSchaller commented Oct 30, 2025

Copy link
Copy Markdown
Contributor Author

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

vfdev-5
vfdev-5 reviewed Oct 31, 2025
View reviewed changes
Comment thread .pre-commit-config.yaml Outdated
Comment thread docs_nnx/api_reference/flax.nnx/index.rst Outdated
@jorisSchaller

jorisSchaller commented Oct 31, 2025

Copy link
Copy Markdown
Contributor Author

@vfdev-5 reverted changes and merged upstream changes

@vfdev-5

vfdev-5 commented Oct 31, 2025

Copy link
Copy Markdown
Collaborator

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

@jorisSchaller

jorisSchaller commented Nov 2, 2025

Copy link
Copy Markdown
Contributor Author

@vfdev-5 Updated

vfdev-5
vfdev-5 reviewed Nov 2, 2025
View reviewed changes
Comment thread tests/pooling_test.py Outdated
@vfdev-5

vfdev-5 commented Nov 2, 2025

Copy link
Copy Markdown
Collaborator

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

@review-notebook-app

review-notebook-app Bot commented Nov 9, 2025

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

@vfdev-5

vfdev-5 commented Nov 10, 2025

Copy link
Copy Markdown
Collaborator

@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 vfdev-5 force-pushed the feat/nnx_pooling branch 2 times, most recently from a107279 to 3146d14 Compare December 5, 2025 08:41
@jorisSchaller @vfdev-5
Refactor pooling operation
004b932 
- Pooling operation moved out of Linen
- Pooling documentation duplicated in nnx
vfdev-5
vfdev-5 approved these changes Dec 5, 2025
View reviewed changes

@vfdev-5 vfdev-5 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@copybara-service copybara-service Bot merged commit de84eb1 into google:main Dec 9, 2025
17 of 18 checks passed
@vfdev-5 vfdev-5 mentioned this pull request Jan 26, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vfdev-5 vfdev-5 vfdev-5 approved these changes

@cgarciae cgarciae cgarciae approved these changes

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

NNX documentation missing pooling operations

3 participants

@jorisSchaller @vfdev-5 @cgarciae