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.

flax/flax/core/nn/__init__.py

Line 38 in 5d896bc

from flax.linen.pooling import (avg_pool as avg_pool, max_pool as max_pool)

But the entire file uses only

flax/flax/linen/pooling.py

Lines 17 to 19 in 5d896bc

	import jax.numpy as jnp
	import numpy as np
	from jax import lax

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!