Merge pull request #3942 from google:fix-nnx-basics-notebook

PiperOrigin-RevId: 638584553
google · May 30, 2024 · 6e63497 · 6e63497
2 parents 144486b + e30754d
commit 6e63497
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 1 deletion.
diff --git a/docs/nnx/index.rst b/docs/nnx/index.rst
@@ -64,7 +64,7 @@ Features
          .. div:: sd-font-normal
 
             NNX makes it very easy integrate objects with regular JAX code
-            via the `Functional API <nnx_basics#the-functional-api>`__.
+            via the `Functional API <nnx_basics.html#the-functional-api>`__.
 
 Basic usage
 ^^^^^^^^^^^^

diff --git a/docs/nnx/nnx_basics.ipynb b/docs/nnx/nnx_basics.ipynb
@@ -17,6 +17,19 @@
     "that have allowed Linen to scale effectively to large codebases."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "skip-execution"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "! pip install -U flax penzai"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,

diff --git a/docs/nnx/nnx_basics.md b/docs/nnx/nnx_basics.md
@@ -20,6 +20,12 @@ like PyTorch.
 Despite its simplified implementation, NNX supports the same powerful design patterns 
 that have allowed Linen to scale effectively to large codebases.
 
+```{code-cell} ipython3
+:tags: [skip-execution]
+
+! pip install -U flax penzai
+```
+
 ```{code-cell} ipython3
 from flax import nnx
 import jax