NNX documentation missing pooling operations

[tux-type]

The NNX API reference does not seem to have entries for pooling operations like nnx.avg_pool or nnx.max_pool.

[cgarciae]

@8bitmp3

[8bitmp3]

thanks @tux-type @cgarciae

[8bitmp3]

Can you help find the source code @cgarciae https://github.com/search?q=repo%3Agoogle%2Fflax+path%3A%2F%5Eflax%5C%2Fnnx%5C%2F%2F+max_pool&type=code

[jorisSchaller]

Hey, I was also looking at the pooling, the flax/core/nn/ __init__.py file import the function from the old linen API. https://github.com/google/flax/blob/5d896bc1a2c68e2099d147cd2bc18ebb6a46a0bd/flax/core/nn/init.py#L38

But the entire file uses only https://github.com/google/flax/blob/5d896bc1a2c68e2099d147cd2bc18ebb6a46a0bd/flax/linen/pooling.py#L17-L19 We should copy it to flax/core/nn/pooling.py in order to have have it in the new nnx API. We should also copy the pooling doc from https://github.com/google/flax/blob/main/docs/api_reference/flax.linen/layers.rst#pooling to a new file called docs_nnx/api_reference/flax.nnx/nn/pooling.rst to have up to date nnx documentation

If you agree with the changes, I can send a PR with them.

[cgarciae]

@jorisSchaller happy to review the PR!