Pooling operation #5057
Pooling operation #5057
Conversation
|
@jorisSchaller thanks for the PR, can you please rebase it on the latest main to avoid unrelated changes? Thanks! |
|
@vfdev-5 let me know if you'd like any additional changes |
.pre-commit-config.yaml
Outdated
| repos: | ||
| - repo: https://github.com/mwouts/jupytext | ||
| rev: v1.13.8 | ||
| rev: v1.18.1 |
There was a problem hiding this comment.
Let's revert the changes in this file
| .. toctree:: | ||
| :maxdepth: 3 | ||
|
|
||
| bridge |
There was a problem hiding this comment.
Same here, let's revert changes in this file.
|
@vfdev-5 reverted changes and merged upstream changes |
|
@jorisSchaller please check this failing test: https://github.com/google/flax/actions/runs/18982441063/job/54231421369?pr=5057 |
|
@vfdev-5 Updated |
|
@jorisSchaller thanks for updates on the PR !
|
- Pooling operation moved out of Linen - Pooling documentation duplicated in nnx
jorisSchaller
force-pushed
the
feat/nnx_pooling
branch
from
2ea485d to
9666481
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
@jorisSchaller please rebase on main and fix the failing check with ruff: https://github.com/google/flax/actions/runs/19210434754/job/54950332158?pr=5057 |
2 participants
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