allow nnx standalone import

cgarciae · cgarciae · commit 6e05d27b9200 · 2026-01-22T15:05:15.000-08:00
diff --git a/docs_nnx/conf.py b/docs_nnx/conf.py
@@ -170,7 +170,7 @@
 doctest_global_setup = """
 import jax
 import jax.numpy as jnp
-from flax import nnx
+import nnx
 
 import logging as slog
 from absl import logging as alog
diff --git a/docs_nnx/examples/gemma.ipynb b/docs_nnx/examples/gemma.ipynb
@@ -122,7 +122,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "import sentencepiece as spm"
    ]
   },
diff --git a/docs_nnx/examples/gemma.md b/docs_nnx/examples/gemma.md
@@ -72,7 +72,7 @@ vocab_path = f'{weights_dir}/tokenizer.model'
 ## Python imports
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 import sentencepiece as spm
 ```
 
diff --git a/docs_nnx/guides/bridge_guide.ipynb b/docs_nnx/guides/bridge_guide.ipynb
@@ -31,7 +31,7 @@
     "import os\n",
     "os.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n",
     "\n",
-    "from flax import nnx\n",
+    "import nnx\n",
     "from flax import linen as nn\n",
     "from flax.nnx import bridge\n",
     "import jax\n",
diff --git a/docs_nnx/guides/bridge_guide.md b/docs_nnx/guides/bridge_guide.md
@@ -31,7 +31,7 @@ And all built-in Linen layers should have equivalent NNX versions! Check out the
 import os
 os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 
-from flax import nnx
+import nnx
 from flax import linen as nn
 from flax.nnx import bridge
 import jax
diff --git a/docs_nnx/guides/checkpointing.ipynb b/docs_nnx/guides/checkpointing.ipynb
@@ -39,7 +39,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "import orbax.checkpoint as ocp\n",
     "import jax\n",
     "from jax import numpy as jnp\n",
diff --git a/docs_nnx/guides/checkpointing.md b/docs_nnx/guides/checkpointing.md
@@ -34,7 +34,7 @@ The Orbax API examples used throughout the guide are for demonstration purposes,
 Import the necessary dependencies, set up a checkpoint directory and an example Flax NNX model - `TwoLayerMLP` - by subclassing [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html).
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 import orbax.checkpoint as ocp
 import jax
 from jax import numpy as jnp
diff --git a/docs_nnx/guides/demo.ipynb b/docs_nnx/guides/demo.ipynb
@@ -17,7 +17,7 @@
    "source": [
     "import jax\n",
     "from jax import numpy as jnp\n",
-    "from flax import nnx"
+    "import nnx"
    ]
   },
   {
diff --git a/docs_nnx/guides/demo.md b/docs_nnx/guides/demo.md
@@ -13,7 +13,7 @@ jupytext:
 ```{code-cell} ipython3
 import jax
 from jax import numpy as jnp
-from flax import nnx
+import nnx
 ```
 
 ### [1] NNX is Pythonic
diff --git a/docs_nnx/guides/filters_guide.ipynb b/docs_nnx/guides/filters_guide.ipynb
@@ -42,7 +42,7 @@
     }
    ],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "\n",
     "class Foo(nnx.Module):\n",
     "  def __init__(self):\n",
diff --git a/docs_nnx/guides/filters_guide.md b/docs_nnx/guides/filters_guide.md
@@ -21,7 +21,7 @@ In this guide you will learn how to:
 In the following example [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) and [`nnx.BatchStat`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.BatchStat) are used as `Filter`s to split the model into two groups: one with the parameters and the other with the batch statistics:
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 
 class Foo(nnx.Module):
   def __init__(self):
diff --git a/docs_nnx/guides/flax_gspmd.ipynb b/docs_nnx/guides/flax_gspmd.ipynb
@@ -44,7 +44,7 @@
     "from jax.sharding import PartitionSpec as P, NamedSharding, AxisType\n",
     "import optax\n",
     "import flax\n",
-    "from flax import nnx\n",
+    "import nnx\n",
     "\n",
     "# Ignore this if you are already running on a TPU or GPU\n",
     "if not jax._src.xla_bridge.backends_are_initialized():\n",
diff --git a/docs_nnx/guides/flax_gspmd.md b/docs_nnx/guides/flax_gspmd.md
@@ -30,7 +30,7 @@ from jax import numpy as jnp
 from jax.sharding import PartitionSpec as P, NamedSharding, AxisType
 import optax
 import flax
-from flax import nnx
+import nnx
 
 # Ignore this if you are already running on a TPU or GPU
 if not jax._src.xla_bridge.backends_are_initialized():
diff --git a/docs_nnx/guides/jax_and_nnx_transforms.rst b/docs_nnx/guides/jax_and_nnx_transforms.rst
@@ -11,7 +11,7 @@ First, let's set up imports and generate some dummy data:
 
 .. testcode:: Flax NNX, JAX
 
-  from flax import nnx
+  import nnx
   import jax
 
   x = jax.random.normal(jax.random.key(0), (1, 2))
diff --git a/docs_nnx/guides/performance.ipynb b/docs_nnx/guides/performance.ipynb
@@ -20,7 +20,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import optax\n",
diff --git a/docs_nnx/guides/performance.md b/docs_nnx/guides/performance.md
@@ -18,7 +18,7 @@ Currently, Flax [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/
 A full resolution _might_ involve developing a C extension (e.g. `flaxlib`) to speed up some of the traversal logic in [`graph.py`](https://github.com/google/flax/blob/main/flax/nnx/graph.py). Before we continue lets an example of a model and a simple training loop:
 
 ```{code-cell}
-from flax import nnx
+import nnx
 import jax
 import jax.numpy as jnp
 import optax
diff --git a/docs_nnx/guides/pytree.ipynb b/docs_nnx/guides/pytree.ipynb
@@ -31,7 +31,7 @@
     }
    ],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
diff --git a/docs_nnx/guides/pytree.md b/docs_nnx/guides/pytree.md
@@ -15,7 +15,7 @@ jupytext:
 Flax NNX's Modules are by default registered as JAX Pytrees, this allows using them throughout most of JAX APIs but in particular JAX transforms and the `jax.tree.*` functions. Thanks to the pytree protocol a simple NNX program might look like this:
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 import jax
 import jax.numpy as jnp
 
diff --git a/docs_nnx/guides/randomness.ipynb b/docs_nnx/guides/randomness.ipynb
@@ -15,7 +15,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "\n",
     "class Model(nnx.Module):\n",
     "  def __init__(self, *, rngs: nnx.Rngs):\n",
diff --git a/docs_nnx/guides/randomness.md b/docs_nnx/guides/randomness.md
@@ -13,7 +13,7 @@ jupytext:
 Flax NNX uses the stateful `nnx.Rngs` class to simplify Jax's handling of random states. For example, the code below uses a `nnx.Rngs` object to define a simple linear model with dropout:
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 
 class Model(nnx.Module):
   def __init__(self, *, rngs: nnx.Rngs):
diff --git a/docs_nnx/guides/surgery.ipynb b/docs_nnx/guides/surgery.ipynb
@@ -33,7 +33,7 @@
     "from jax.sharding import PartitionSpec, Mesh, NamedSharding\n",
     "from jax.experimental import mesh_utils\n",
     "import flax\n",
-    "from flax import nnx\n",
+    "import nnx\n",
     "import flax.traverse_util\n",
     "import numpy as np\n",
     "import orbax.checkpoint as orbax\n",
diff --git a/docs_nnx/guides/surgery.md b/docs_nnx/guides/surgery.md
@@ -31,7 +31,7 @@ from jax import lax, numpy as jnp, tree_util as jtu
 from jax.sharding import PartitionSpec, Mesh, NamedSharding
 from jax.experimental import mesh_utils
 import flax
-from flax import nnx
+import nnx
 import flax.traverse_util
 import numpy as np
 import orbax.checkpoint as orbax
diff --git a/docs_nnx/guides/transforms.ipynb b/docs_nnx/guides/transforms.ipynb
@@ -31,7 +31,7 @@
    "source": [
     "import jax\n",
     "from jax import numpy as jnp, random\n",
-    "from flax import nnx"
+    "import nnx"
    ]
   },
   {
diff --git a/docs_nnx/guides/transforms.md b/docs_nnx/guides/transforms.md
@@ -29,7 +29,7 @@ Flax NNX objects, they allow Python's reference semantics to be preserved for th
 ```{code-cell} ipython3
 import jax
 from jax import numpy as jnp, random
-from flax import nnx
+import nnx
 ```
 
 Throughout this guide, `nnx.vmap` is used as a case study to demonstrate how Flax NNX transforms work. However, the principles
diff --git a/docs_nnx/hijax/hijax.ipynb b/docs_nnx/hijax/hijax.ipynb
@@ -15,7 +15,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import optax\n",
diff --git a/docs_nnx/hijax/hijax.md b/docs_nnx/hijax/hijax.md
@@ -11,7 +11,7 @@ jupytext:
 # Hijax
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 import jax
 import jax.numpy as jnp
 import optax
diff --git a/docs_nnx/hijax/index.rst b/docs_nnx/hijax/index.rst
@@ -13,7 +13,7 @@ Basic usage
 
 .. testcode::
 
-  from flax import nnx
+  import nnx
   import optax
 
   nnx.var_defaults(hijax=True)
diff --git a/docs_nnx/index.rst b/docs_nnx/index.rst
@@ -89,7 +89,7 @@ Basic usage
 
 .. testcode::
 
-   from flax import nnx
+   import nnx
    import optax
 
 
diff --git a/docs_nnx/key_concepts.ipynb b/docs_nnx/key_concepts.ipynb
@@ -24,7 +24,7 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import flax\n",
-    "from flax import nnx\n",
+    "import nnx\n",
     "from functools import partial\n",
     "\n",
     "# For simulating multi-device environment\n",
diff --git a/docs_nnx/key_concepts.md b/docs_nnx/key_concepts.md
@@ -11,7 +11,7 @@ This means a **basic understanding on JAX helps you to use Flax well**. You woul
 import jax
 import jax.numpy as jnp
 import flax
-from flax import nnx
+import nnx
 from functools import partial
 
 # For simulating multi-device environment
diff --git a/docs_nnx/migrating/convert_pytorch_to_flax.rst b/docs_nnx/migrating/convert_pytorch_to_flax.rst
@@ -6,7 +6,7 @@ Convert PyTorch models to Flax
   import numpy as np
   import jax
   from jax import random, numpy as jnp
-  from flax import nnx
+  import nnx
 
   import torch
 
diff --git a/docs_nnx/migrating/haiku_to_flax.rst b/docs_nnx/migrating/haiku_to_flax.rst
@@ -66,7 +66,7 @@ There are two fundamental differences between Haiku and Flax ``Module`` objects:
 
   ---
 
-  from flax import nnx
+  import nnx
 
   class Block(nnx.Module):
     def __init__(self, in_features: int , out_features: int, rngs: nnx.Rngs):
diff --git a/docs_nnx/migrating/linen_to_nnx.rst b/docs_nnx/migrating/linen_to_nnx.rst
@@ -57,7 +57,7 @@ There are two fundamental differences between Flax Linen and Flax NNX ``Module``
 
   ---
 
-  from flax import nnx
+  import nnx
 
   class Block(nnx.Module):
     def __init__(self, in_features: int , out_features: int, rngs: nnx.Rngs):
diff --git a/docs_nnx/migrating/nnx_010_to_nnx_011.rst b/docs_nnx/migrating/nnx_010_to_nnx_011.rst
@@ -145,7 +145,7 @@ the new changes:
     .. code-block:: python
       :emphasize-lines: 17, 26
 
-      from flax import nnx
+      import nnx
       import optax
 
 
@@ -180,7 +180,7 @@ the new changes:
     .. code-block:: python
       :emphasize-lines: 17, 26
 
-      from flax import nnx
+      import nnx
       import optax
 
 
diff --git a/docs_nnx/mnist_tutorial.ipynb b/docs_nnx/mnist_tutorial.ipynb
@@ -124,7 +124,7 @@
     }
    ],
    "source": [
-    "from flax import nnx  # The Flax NNX API.\n",
+    "import nnx  # The Flax NNX API.\n",
     "from functools import partial\n",
     "from typing import Optional\n",
     "\n",
diff --git a/docs_nnx/mnist_tutorial.md b/docs_nnx/mnist_tutorial.md
@@ -72,7 +72,7 @@ test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1)
 Create a CNN for classification with Flax NNX by subclassing `nnx.Module`:
 
 ```{code-cell} ipython3
-from flax import nnx  # The Flax NNX API.
+import nnx  # The Flax NNX API.
 from functools import partial
 from typing import Optional
 
diff --git a/docs_nnx/nnx_basics.ipynb b/docs_nnx/nnx_basics.ipynb
@@ -30,7 +30,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flax import nnx\n",
+    "import nnx\n",
     "import jax\n",
     "import jax.numpy as jnp"
    ]
@@ -742,7 +742,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,
diff --git a/docs_nnx/nnx_basics.md b/docs_nnx/nnx_basics.md
@@ -21,7 +21,7 @@ To begin, install Flax with `pip` and import necessary dependencies:
 ```
 
 ```{code-cell} ipython3
-from flax import nnx
+import nnx
 import jax
 import jax.numpy as jnp
 ```
diff --git a/docs_nnx/why.rst b/docs_nnx/why.rst
@@ -19,7 +19,7 @@ complexity and obscured the core ideas of `transformations (transforms) <https:/
 
     import jax
     from jax import random, numpy as jnp
-    from flax import nnx
+    import nnx
     import flax.linen as nn
 
 Introducing Flax NNX
@@ -37,7 +37,7 @@ Here is an example of a simple Flax NNX program that illustrates many of the poi
 
 .. testcode:: NNX
 
-  from flax import nnx
+  import nnx
   import optax
 
 
diff --git a/nnx.py b/nnx.py
@@ -31,4 +31,4 @@
 from flax import nnx as _nnx
 
 # Re-export the module's metadata
-__all__: list[str] = _nnx.__all__ if hasattr(_nnx, '__all__') else []
+__all__: list[str] = _nnx.__all__ if hasattr(_nnx, '__all__') else []
diff --git a/pyproject.toml b/pyproject.toml
@@ -113,6 +113,9 @@ version = {attr = "flax.version.__version__"}
 [tool.setuptools.packages.find]
 include = ["flax*"]
 
+[tool.setuptools]
+py-modules = ["nnx"]
+
 [tool.setuptools.package-data]
 flax = ["*py.typed"]
 

Original file line number	Diff line number	Diff line change
`@@ -122,7 +122,7 @@`
`122`	`122`	`"metadata": {},`
`123`	`123`	`"outputs": [],`
`124`	`124`	`"source": [`
`125`		`- "from flax import nnx\n",`
	`125`	`+ "import nnx\n",`
`126`	`126`	`"import sentencepiece as spm"`
`127`	`127`	`]`
`128`	`128`	`},`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@`
`17`	`17`	`"source": [`
`18`	`18`	`"import jax\n",`
`19`	`19`	`"from jax import numpy as jnp\n",`
`20`		`- "from flax import nnx"`
	`20`	`+ "import nnx"`
`21`	`21`	`]`
`22`	`22`	`},`
`23`	`23`	`{`
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@`
`42`	`42`	`}`
`43`	`43`	`],`
`44`	`44`	`"source": [`
`45`		`- "from flax import nnx\n",`
	`45`	`+ "import nnx\n",`
`46`	`46`	`"\n",`
`47`	`47`	`"class Foo(nnx.Module):\n",`
`48`	`48`	`" def __init__(self):\n",`