Pooling functions
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx

.. autofunction:: avg_pool
.. autofunction:: max_pool
.. autofunction:: min_pool
.. autofunction:: pool