Merge pull request #3503 from google:nnx-fix-why

PiperOrigin-RevId: 584894048
google · Nov 23, 2023 · 97b6167 · 97b6167
2 parents 70214f4 + 4308746
commit 97b6167
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/flax/experimental/nnx/docs/why.ipynb b/flax/experimental/nnx/docs/why.ipynb
@@ -34,7 +34,8 @@
    "source": [
     "from functools import partial\n",
     "import jax\n",
-    "from jax import random, numpy as jnp"
+    "from jax import random, numpy as jnp\n",
+    "from flax.experimental import nnx"
    ]
   },
   {

diff --git a/flax/experimental/nnx/docs/why.md b/flax/experimental/nnx/docs/why.md
@@ -33,6 +33,7 @@ We'd love to hear from any of our users about their thoughts on these ideas.
 from functools import partial
 import jax
 from jax import random, numpy as jnp
+from flax.experimental import nnx
 ```
 
 ### NNX is Pythonic