Merge pull request #3936 from google:nnx-stabilize

PiperOrigin-RevId: 636951986

Author: Flax Authors

CHANGELOG.md

@@ -82,7 +82,7 @@ vNext
 
 0.8.0
 -----
-- Added [NNX](https://github.com/google/flax/tree/main/flax/experimental/nnx#nnx), a neural network library for JAX that provides a simple yet powerful module system that adheres to standard Python semantics. Its aim is to combine the robustness of Linen with a simplified, Pythonic API akin to that of PyTorch.
+- Added [NNX](https://github.com/google/flax/tree/main/flax/nnx#nnx), a neural network library for JAX that provides a simple yet powerful module system that adheres to standard Python semantics. Its aim is to combine the robustness of Linen with a simplified, Pythonic API akin to that of PyTorch.
 - Added `nn.compact_name_scope` decorator that enables methods to act as compact name scopes as with regular Haiku methods. This makes porting Haiku code easier.
 - Add copy() method to Module.  This is a user-friendly version of the internal clone() method with better
   defaults for common use cases.

README.md

@@ -12,6 +12,8 @@
 | [**What does Flax look like?**](#what-does-flax-look-like)
 | [**Documentation**](https://flax.readthedocs.io/)
 
+**📣 NEW**: Check out the [**NNX**](https://flax.readthedocs.io/en/latest/nnx/index.html) API!
+
 This README is a very short intro. **To learn everything you need to know about Flax, refer to our [full documentation](https://flax.readthedocs.io/).**
 
 Flax was originally started by engineers and researchers within the Brain Team in Google Research (in close collaboration with the JAX team), and is now developed jointly with the open source community.

docs/api_reference/flax.experimental.nnx/nn/stochastic.rst

@@ -1,8 +1,8 @@
 graph
 ------------------------
 
-.. automodule:: flax.experimental.nnx
-.. currentmodule:: flax.experimental.nnx
+.. automodule:: flax.nnx
+.. currentmodule:: flax.nnx
 
 
 .. autofunction:: split

docs/api_reference/flax.experimental.nnx/training/optimizer.rst

@@ -1,8 +1,8 @@
 helpers
 ------------------------
 
-.. automodule:: flax.experimental.nnx
-.. currentmodule:: flax.experimental.nnx
+.. automodule:: flax.nnx
+.. currentmodule:: flax.nnx
 
 .. autoclass:: Dict
    :members:

docs/api_reference/flax.experimental.nnx/visualization.rst

@@ -1,7 +1,7 @@
-flax.experimental.nnx
+flax.nnx
 ------------------------
 
-Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/experimental/nnx/index.html>`__ for more details.
+Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/nnx/index.html>`__ for more details.
 
 .. toctree::
   :maxdepth: 3

docs/api_reference/flax.experimental.nnx/graph.rst renamed to docs/api_reference/flax.nnx/graph.rst

@@ -1,8 +1,8 @@
 module
 ------------------------
 
-.. automodule:: flax.experimental.nnx
-.. currentmodule:: flax.experimental.nnx
+.. automodule:: flax.nnx
+.. currentmodule:: flax.nnx
 
 .. autoclass:: Module
    :members:

docs/api_reference/flax.experimental.nnx/helpers.rst renamed to docs/api_reference/flax.nnx/helpers.rst

@@ -1,8 +1,8 @@
 Activation functions
 ------------------------
 
-.. automodule:: flax.experimental.nnx
-.. currentmodule:: flax.experimental.nnx
+.. automodule:: flax.nnx
+.. currentmodule:: flax.nnx
 
 .. autofunction:: celu
 .. autofunction:: elu