[nnx] fix nnx_basics dependencies

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
 ^^^^^^^^^^^^

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,

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