google / google/flax

feat(nnx): Missing object-oriented pooling layers in NNX

Open
#5,202 2 comments 1 reaction 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.