feat(nnx): Missing object-oriented pooling layers in NNX
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
### Problem Description
Currently, `flax.nnx` lacks native object-oriented pooling modules such as `MaxPool`, `AvgPool`, and `GlobalAveragePool`. Users migrating from frameworks like Keras or PyTorch—or even transitioning from `flax.linen`—are forced to mix functional API calls within the object-oriented NNX structure. This creates an inconsistent developer experience and requires manual boilerplate for common operations like Global Average Pooling.
### Proposed Feature
Introduce a dedicated pooling module suite within `nnx` that mirrors the ergonomic design of other NNX layers. This includes:
* **Subsampling Modules**: `MaxPool`, `AvgPool`, and `MinPool`.
* **Global Pooling**: A dedicated `GlobalAveragePool` module to replace manual `jnp.mean` calls.
### Implementation Status
I have already implemented these modules and exposed them in the `nnx` namespace.
**See Pull Request: #5201**
### Justification & Benefits
1. **API Consistency**: Maintains the OO-flow of NNX without jumping back into `linen.functional`.
2. **Framework Parity**: Lowers the barrier for users migrating from Keras/PyTorch.
3. **Readability**: Simplifies model definitions, especially for standard CNN architectures.
Contributor guide
Assessment
This issue has not been assessed yet.