Repository: google/flax

Issue:
Support for RNN (NNX)
seems like that RNN family is not currently supported in the NNX api

Discussion:
@zinccat: started working on porting: https://github.com/zinccat/rnn_nnx@cgarciae: Hey @zinccat, if you want to send PRs happy to review.
We haven't ported any of the cells or the `RNN` class itself.

Solution patch:
diff --git a/.github/workflows/flax_test.yml b/.github/workflows/flax_test.yml
index 1420f39a4..db737ed27 100644
--- a/.github/workflows/flax_test.yml
+++ b/.github/workflows/flax_test.yml
@@ -75,7 +75,7 @@ jobs:
         uv-version: "0.3.0"
     - name: Install standalone dependencies only
       run: |
-        uv sync --locked --extra all
+        uv sync --extra all
     - name: Test importing Flax
       run: |
         uv run python -c "import flax"
@@ -113,7 +113,7 @@ jobs:
 
     - name: Install dependencies
       run: |
-        uv sync --locked --extra all --extra testing --extra docs
+        uv sync --extra all --extra testing --extra docs
         uv pip install ./flaxlib
     - name: Install JAX
       run: |
diff --git a/dev/.devcontainer/Dockerfile b/dev/.devcontainer/Dockerfile
index 11fb57e0b..6f48300f7 100644
--- a/dev/.devcontainer/Dockerfile
+++ b/dev/.devcontainer/Dockerfile
@@ -3,7 +3,7 @@
 # Licensed under the MIT License. See https://go.microsoft.com/fwlink/?linkid=2090316 for license information.
 #-------------------------------------------------------------------------------------------------------------
 
-FROM python:3.7
+FROM python:3.12
 
 # Avoid warnings by switching to noninteractive
 ENV DEBIAN_FRONTEND=noninteractive
@@ -52,7 +52,7 @@ RUN apt-get update \
 RUN pip install numpy jaxlib tensorflow tensorflow-datasets matplotlib msgpack \
     jupyter pytest pytest-xdist \
     twine \
-    sphinx sphinx_rtd_theme ipykernel nbsphinx recommonmark sklearn
+    sphinx sphinx_rtd_theme ipykernel nbsphinx recommonmark scikit-learn
 
 # Switch back to dialog for any ad-hoc use of apt-get
 ENV DEBIAN_FRONTEND=dialog
diff --git a/dev/.devcontainer/devcontainer.json b/dev/.devcontainer/devcontainer.json
index 32b87efa7..2b6cde283 100644
--- a/dev/.devcontainer/devcontainer.json
+++ b/dev/.devcontainer/devcontainer.json
@@ -22,7 +22,7 @@
 	// Uncomment the next line to run commands after the container is created.
 	// "postCreateCommand": "pip install -r requirements.txt",
 	"runArgs": ["-v", "${env:HOME}${env:USERPROFILE}/.ssh:/root/.ssh-localhost:ro"],
-	"postCreateCommand": "sudo cp -r /root/.ssh-localhost ~/.ssh && sudo chown -R $(id -u):$(id -g) ~/.ssh && chmod 700 ~/.ssh && chmod 600 ~/.ssh/* && pip install -e ./flax",
+	"postCreateCommand": "sudo cp -r /root/.ssh-localhost ~/.ssh && sudo chown -R $(id -u):$(id -g) ~/.ssh && chmod 700 ~/.ssh && chmod 600 ~/.ssh/* && pip install -e .",
 
 	// Uncomment the next line to have VS Code connect as an existing non-root user in the container.
 	// On Linux, by default, the container user's UID/GID will be updated to match your local user. See
diff --git a/dev/README.md b/dev/README.md
index 5d8f081bb..56f3484b5 100644
--- a/dev/README.md
+++ b/dev/README.md
@@ -13,4 +13,7 @@ the environments used by contributors and maintainers.
 5. Re-open the folder workspace using the remote containers extension.
    VSCode should recommend this action in a popup. Alternatively,
    use the green button in the bottom left container to control the
-   remote extension.
\ No newline at end of file
+   remote extension.
+
+## Troubleshoot:
+If you have the following error `~/.docker/buildx/current: permission denied`, try running `sudo chown -R $(whoami) ~/.docker`
\ No newline at end of file
diff --git a/docs_nnx/conf.py b/docs_nnx/conf.py
index c3deabe87..4a6075455 100644
--- a/docs_nnx/conf.py
+++ b/docs_nnx/conf.py
@@ -148,9 +148,6 @@
   'transfer_learning.ipynb',  # <-- transformers requires flax<=0.7.0
   'flax/nnx',  # exclude nnx
   'guides/demo.ipynb',  # TODO(cgarciae): broken, remove or update
-  'guides/why.ipynb',  # TODO(cgarciae): broken, remove in favor on the new guide
-  'guides/flax_gspmd.ipynb',  # TODO(IvyZX): broken, needs to be updated
-  'guides/surgery.ipynb',  # TODO(IvyZX): broken, needs to be updated
 ]
 # raise exceptions on execution so CI can catch errors
 nb_execution_allow_errors = False
diff --git a/docs_nnx/guides/checkpointing.ipynb b/docs_nnx/guides/checkpointing.ipynb
index 100dcffc3..449f8a775 100644
--- a/docs_nnx/guides/checkpointing.ipynb
+++ b/docs_nnx/guides/checkpointing.ipynb
@@ -6,11 +6,22 @@
    "source": [
     "# Save and load checkpoints\n",
     "\n",
-    "Flax does not actively maintain a library for saving and loading model checkpoints to disk. We recommend you to use external libraries like [Orbax](https://orbax.readthedocs.io/en/latest/index.html) to do it.\n",
+    "This guide demonstrates how to save and load Flax NNX model checkpoints with [Orbax](https://orbax.readthedocs.io/).\n",
     "\n",
-    "This guide will cover a few user-side scenarios of saving and loading Flax NNX model checkpoints using Orbax. The Orbax API we use here are for demonstration purposes only - please check out the [Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for the most up-to-date recommended API.\n",
+    "> **Note:** The Flax team does not actively maintain a library for saving and loading model checkpoints to disk. Therefore, it is recommended you use external libraries like [Orbax](https://orbax.readthedocs.io/en/latest/index.html) to do it.\n",
     "\n",
-    "> Note: Flax's legacy `flax.training.checkpoints` package is deprecated. Its documentation resides in the [legacy documentation site](https://flax-linen.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html)."
+    "In this guide you will learn how to:\n",
+    "\n",
+    "* Save checkpoints.\n",
+    "* Restore checkpoints.\n",
+    "* Restore checkpoints if checkpoint structures differ. \n",
+    "* Perform multi-process checkpointing. \n",
+    "\n",
+    "The Orbax API examples used throughout the guide are for demonstration purposes, and for the most up-to-date recommended APIs refer to the [Orbax website](https://orbax.readthedocs.io/).\n",
+    "\n",
+    "> **Note:** The Flax team recommends using [Orbax](https://orbax.readthedocs.io/en/latest/index.html) for saving and loading checkpoints to disk, as we do not actively maintain a library for these functionalities.\n",
+    "\n",
+    "> **Note:** If you are looking for Flax Linen's legacy `flax.training.checkpoints` package, it was deprecated in 2023 in favor of Orbax. The documentation resides on the [Flax Linen site](https://flax-linen.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html)."
    ]
   },
   {
@@ -19,7 +30,7 @@
    "source": [
     "### Setup\n",
     "\n",
-    "We first set up a checkpoint directory and an example NNX model."
+    "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)."
    ]
   },
   {
@@ -52,7 +63,7 @@
     "    x = self.linear1(x)\n",
     "    return self.linear2(x)\n",
     "\n",
-    "# Create this model and show we can run it\n",
+    "# Instantiate the model and show we can run it.\n",
     "model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\n",
     "x = jax.random.normal(jax.random.key(42), (3, 4))\n",
     "assert model(x).shape == (3, 4)"
@@ -64,9 +75,9 @@
    "source": [
     "## Save checkpoints\n",
     "\n",
-    "JAX checkpointing libraries save [Pytrees](https://jax.readthedocs.io/en/latest/pytrees.html), which is a pure, possibly nested container of JAX arrays (or, \"tensors\" as some other frameworks would put it). In the context of machine learning, the checkpoint is usually a pytree of model parameters and other data such as optimizer states.\n",
+    "JAX checkpointing libraries, such as Orbax, can save and load any given JAX [pytree](https://jax.readthedocs.io/en/latest/pytrees.html), which is a pure, possibly nested container of [`jax.Array`s)](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) (or, \"tensors\" as some other frameworks would put it). In the context of machine learning, the checkpoint is usually a pytree of model parameters and other data, such as optimizer states.\n",
     "\n",
-    "In Flax NNX, you can obtain such a pytree from an [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) by calling [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split) and pick up the returned [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State)."
+    "In Flax NNX, you can obtain such a pytree from an [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) by calling [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), and picking up the returned [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State)."
    ]
   },
   {
@@ -111,13 +122,20 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Restore checkpoints\n",
+    "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_5ecb69e45849446495cac4e3367e5c4f\" ></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5ecb69e45849446495cac4e3367e5c4f\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWgtT27oS/itqOnNILsRJnBcJj7lOyIsWKIQWyjlncmVbdkQc29hKQjjDf78rOe+YFFoK5TQw0xBppV3tS9+uuuuzoUX2JeYR4muOS1qe4zD0D3IdnzLq2EXkEQsz2ic7yHBsFjdwl1rDIuo6tuO7WIPxQZsyEhdfisj1YMSiPouLreNs6MKo7dgwrGKtY3pOz9bjmmM5XjFYuoNG31QLCGA/qrN2ERmUAZnNiM12kIt1ndpm3CIGKyJZa3MmNom3CTXbMJKSsnwbm2EKMk+Wjf6I96lPVWpRBpLjHnMmtHFqM4/aPtXiPr0jwexI3PvdRKCe3Yl64l7PBp4ejPmaR12G+Pn2NrDrWlTDXGMJR2OEn94juLuxH43G9vZBocDPZ0gnhu2jPcTa1JdMws5A28eOTqIxqe34TBLzcDTCUMslNj+yovFd+aI//w6bqWNbtwhM2z3L2gk4SCBm03FsGI0OHK8TQ7MyOBcwxKfmhhnV+KBLPMPxutjWiGQ7g2hM2BcYRJdmUDxYtIvScgz2oQaKLkgtWcQ2WRvt7aEkJ1kpukdYz7NB74hYPpkK1u7ZXLLFrf02NRiXTxDwP+7h9wEOUfAqW3cGkkduesRnik27wlxVD3dJNNBJjO+xs8TI7fntQI07IWccs9gLjrHilI+XgUsRGJI5pmkFUdkSkQPe6vK9+Aix2BYifXDwkSW5dOK71CFDrvSIF+ECjYglzcK+/xGCc7RvNDLZs9UFN4yMmd/HQJ/g/sLH93cTYQGg0z4SG+5F5tNHBDGswknJ7V4kGYHQ9dgyiWODiKAMG6ZWBUO4BqJ8zfjsEQjGII2JPNLCquqRvvAfkVbe57ZlnEzCqUYEmtPtwsIZCix++OEXSHDRdli02Hb6xIuF0M+TtwZtYrfIrQsmJ7pYKhmOpWMVTmDD0Ypt7Ef3LawSa39+phWcM2CntYnWIXoshv4TQw9zBQ2CW+ozFEn4MYwphc+Tmzm7RTqVTWU5AWiXeB7RWy5kXdIGWYg3S5jnv6OkL9RbRJRhyHJ88ZzoD1wXoMuwEwK1Tn1gOhxfC4uEaB8JDRWLKoGMQ2ak0sTPTii/4GqIp/jdMLpCwOYTXtQWF4ZqOfyOeZCnsPQyZx17HZ9gE7zYXl79TJaeyMCXhi8a089JKO64Itr4S86q2sZrije/6EEhcy8gJLcjZ9zzfG5A14GbnnghfKn/fGxFKAhGcZGb/Id8/Hm4To/HyC1b5iJRv2VQz2ctx25x9w8JrVWhJMlZHk2hpkI/LH5g8UUR+am62DMBmwViiIC+/0FukA/dodpjDEBRWAKaToc5bQRFFqhAkQB4w4n/IqmMHlnAyxtHGLyCYgs1h13VsXx00mP8vDoqByvh0x1CYMQHRO0A9g0ybxfuvTbkcECmNoPlFPtEnyDm9yTJf3eW3TxYLSBtUiqQ7uIpg/gIOUV4upuulAbYb2lwD4BiJ+uxweZuj3GeXsVzYc08y1nVoz72ovG4jhmOYxsMK0BTbHaYM+FI0MP22JvFtijlIwIaA5gfd3rsaUeZSACGoUR/Ny+JYIne0a7reAzbS3urntMBMMBHpsno29qdWTajz7GZ7yWOq0AwvaUBKtc9Yo9Ena97YM95wmfDIpPQGV2kc6GqYUuLQnEGlUHKvRVgUvIZ5usn8v40SVSAQpB3A0l0h8HZuRSzyrvpYcsGrN2CKtWgt7DJXJhsizCBwgJzXDTAng2B1xon9rEtDANrqXQIoQvI/J9JpeqNClOevUZKGg3Fk5JIq9NyuCiKWezFTQ/rFMwWTaWzOjG3kAMebRKUBOlyWnsr8HBA1DxfiCE00vKSKEuJ9dkUP2fy8WHuJYF9QUkWdiFFfRvmPT2HP8whAHeCxxwEf4DmOeQIYzE66DxBC3D4tHQLU8vD5M9jsy200OCR5ss+9EiBF7U8MUUL1oEI7IeXveh5Hyn+nFFXG/IxXB+losfo8jHMvnWo5+hw8aYDeqd4Hh5Khud0o7qj9Xg1LPHk6kt9bPWIH43FJN/pkqhIubxdwT+lAB7xVsUjAVJkA8VQbNIc8tuEMN5BIgNUbjab/DRNPsb7QWJS8oiocJtDW4v+778jUKaRcfJ/OkCbLYlt3hGzRmODUTMyw5sNvqcVUc+zohwtFPl8YuAYhryjAh7JZbb0ZKF2ZColRfw0ThXFEX+Vzgbwb72qKBVl1U+pqyhmx/mgNyql8uCropx/LR8qR41SWamat436xzbzS0eUmOnqwaX8sZH72m+6PfrpKHueOrxsnH056l8c3bFPw2q1vHlhds5p6SDZpgenvcOKXrtO1tWE0W/o7s2HXPvmgtLT3pFda9eNz0z5nCsdexml2rA7lZz2udezN8+yN5rfGfSNqpW4uTUrzrapHg5q26m6krCVs+xHzztMnW2ad8kzPakcGinzOF8e1K5lM+kMe2f5fLeSyg3ql4UT03TJeWeYIQ31Lqup3kmNYcU8bRwPDrA/9E97jcblRaU6UD6duo2v+udEYtPMn+cv0yxpfPh0o/SzsOdH5TivHA2Urnl31tzsXTVJ5fJWNnLa3XHmrD7M9krKh7vStVt107R+Wq4kr3qfMs28bZQ+VurVo65CN7f7Fbltp9r5TfXL4PJ6UPf6B7XPZfvaqFRMtnmiXVlWPlsoHw5K2+1C5uio1kzXrhSz28hel04L7LxG6oVKqdSopQ/MzFniqzZUlRrY9MuHhHJawwo5KltK/a5yYl4xM1f6ZJ6cNA5KHXqaJdXSZblU1WjSbXuOa4NvuFeVg9RdqtM0ygZrDz/YdR1X/bqRPO7WKse5kq7cfPniYuY3r7q6jmlBNu4Kmc/0+ibndr3cifO13KRerds/rKWbF810tSJrpVPjfLNuOW4tU/UHWWze5LbpFWkeW+6FXao3iH7kkd7FTa3cTV1UvU6zeZuVcxcX/kABiWJIdINZdEO49QbHIf+DfybRj3XHBQw2DUnRw5YkaQXFVhCzf8Neq7uCbdFUFTA5QPCwN7iHraFoAKTnW94QgucOD18gGwFtPuZDeuBb8MqBw208wJQhG/epiZnjSbCzqzrY06WBRxk5hyI7Ot0LDjvaa9pXBfAXjcyUFbyjClzOaZdA/REdt9yX1nmkCzXC0tL7LSQnk0mB8SD5AtyLigI5nO9M7RCZCsdbA+MMxpvQEfQeVTG1ILExB3HidyKzAa6zAR9DNqagM4J1Xvpszupu1B3+Rl+Yl1HjxvB8j20RkUb2d4N7fZfabm9000TETa46t5HQTUaXPkwGFz4IIRbP852/aSP7f1iMiw0Uq+nmJheqhci+YeFbybZveU3DiEVVabxn8NHkw9F/dhOjU81utlS3RVbNL0zONnPhMO9v5fxOUDGkgi8jAYroYe0/q7JDz7gxLWM3kGOXuRvvbTwxjkVnP7aBJjX2XuTPjdFZN/6OIHHv7kVmivAi+uOm57CdGbJgYActNQ3gmheeDJmyDX+P1Lw/b8afazeVYv9FjPYqkdHHAKFAEn8hNL6MxoMQeSXv+XODK//RXjQmf0VvEvZ+SfdY8ZwVCXGQ6dx79KAXhdaVP+xRn7CHu5OhkIy7SvhniOrglTEIIDGBgqgOEVyIGsz+Ya4It3BhXzAwJO5AT4wOseYZQ2TOAyP7CNDKAbgEr8MwE30ylMtuI8dAic8+8fwE7Q/vQHtmIjlIJrj2E2oPjJZotYhOmfD7lsTH40mpICXj7jAd5+UvSDkM6MFaiam13OGCdCs+tpaNOOkqgOxrS69W4qqQE22DvRDXGL16R/aPwYiP0PDWL2wFccinmkEselY7BB+x1UBhJnO9ced/2fv/8S7/TezYIZ5NrDV6fCX0GKj/Cf4zXrBGkGsE+W9HkGNffxqymFu1RpFv5iJ9fWv/EJIMv5+ovheZ/LfIbaKmUljN5PF2MpPZTmO9kMxkk9ntgpohqZyxENJh/50ysn8NkSxexpBhOZil5WhmC2Vi3z7D8utdqJCGpqeS6XQ+l8mkM9taDhs52cjn5TTJpnAmk4vs/+zUH5Y3w3XB09wD+ljGlt9SzffD+5XY5GVj52nof37Zz82VX6gP1ym9EzuiNtUhohC10fQtmT8yf5f//maKfsZ6680Vua+Fl8M+7t/qXfsyTx7PWKQGIsm/0cOU/DgryeuHqXVrIdR7ntKYktcPU+u2wm/RVpC/47lCXj9MvUGY8wtYev0wNafRR1dM8vph6q3e/+uHqX8LenxKoS2vH6bWCPJ3QpBPfqqQ1w9Tb/cifX1r/+yHKVnOYbWQklM5OZ3JGakCTuVVdTtfkFN6IZnXf4mHqWxaxnpSy+t5PZXJZ7NqKl/Q8oR/z+XknL5+mPq1Hqbk73svkdcPU29V0b/pw9Sr4uVVD1NvTIkv8+ARqrDndM8JpU77+/8HnO2ZvA==</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5ecb69e45849446495cac4e3367e5c4f\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
     "\n",
-    "Note that you saved the checkpoint as a class of `nnx.State`, which is also nested with `nnx.VariableState` and `nnx.Param` Python classes. At restoration time, you need to have these classes ready in your runtime, and tell the checkpointing library to restore your pytree back to that structure.\n",
     "\n",
-    "You can achive this by creating an abstract NNX model (without allocating any memory for arrays) and show its abstract variable state to the checkpointing library.\n",
+    "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet62ziy4H8/BduTbkqxJJO6y7I9Xy59yekknUm6p0+vjz+HIiGbCUWqScqR26P/O++x+wD7Cvso8yRbBYAkAIKU7CRn9vs66UtEoKpQKBSAqsLt2POvjSS9CcjJvucny8C5OTLCKCT7hu+d7M+j+MIjcxLHxLuYu55t9XqjYb/f64/doTMfduejUbdHBrbT7w/3T4+TpRPC/5HeaceJY+fm2v/jwo3C1PFDEhu3xocrPyVtgHMJFhQvnGBqbAwdcMcP5xGgzCGlPXcWfgC8LaIwothTw42CKD4y/uLQP1Nj4cSXftieRWkaLY4Mq9MdkMVULnEZk/ri/HC5Ss/SmyUIJHbCS7J/Dixckzj1XSdoO4F/GQIXvucFQGnuBykBHi6BWgL5pGE3jQiK8tObhtUZNO9c2NFVdE0FVSZ9N3rhajEjMRAMo7RxNI/cVdIEsrMo9kjcjh3PXyVHRm+5/jiS7DdlGslnbTKif6a8uCPDXq6NJAp8r8iqKbWTACSJE1Vf6lqPspD6S8CRFHlqLKPET/0Ims2ZAQ+rFNJmjvv+Mo5WodfmLNOCdAzPAoAFKo7n+eEl0yv3Csn6IbRQm1yTME2ywj74Xnp1BK2XtpE5yJoayNk8iD4cGdd+4s9QccrV+qPthx5ZQ8mWZdXXchatd6xltG4nV46HRVv0H6wWrVCLJ3QhgVddX6Gcr0kNW27gu+89J3Xu0mJB5KBELxYkSZxLImhP1qM3x4dsLDlOY0ISN1qSdrwK21ckhrTEjf1lalDdNJ3lEnhwUAKHkZuStJ0AjrMwT/fwDxSbpEbGhXFiNBpN4+TUuN0zDPh3vgpdRDU8kpDYh17+B/kFpDFuYKcBAMOISbqKQ4OmPkI6nXkcLRpOGs0AqGU0FpTgouNGHnmFonyUNqxmcwrYm73qYr4DMaS9blEQY3V2k5IE+LxXeRmROdJGKiH5YPCCKK0GJd+ZreYwtnMUXkGGs43rZ+F/C8+0mDtxzFgOSGo8QWVaOMvX3z9+Cpo5VWtzSdInoIx+uIpWCQVuXDvBirSYGgImomU1RIozJyEXtDe0jGg+T0jK+PDnBkM1jk8MK8MwBHiojjXlqQyzSNkYJEiIQOT0xLAriIicdQISXqZXRtvolkjbHZl4RoyJ2E2TnCIr8qHR0JO2m1MdHy+c9KoDcgeZ5cSaJS6Kcr42bM6P0NKxUqGzoojzM+scmbKBBUauaRxw8kYVknFg2BxRbB1W2GVdYfZ9C7P1hc3qCuvet7CuWhjX/7O4ZVy2jNm5vtPehGA6uY9iN/HDq9cEqDd4ee/JDR3z/87VPvCXLxzQ7dj58MIP2d/4zUl87ywztcypJ6kD89kbtE88XkQDMFJnVWgwavZXfvKdH8LE0KBZ//gHUyGYqhrrpnGICMaxYZN2v8DLK7jONEvR5hyA0krAMkNiDymxhzkM/qEAQXTZKJd6wLF/j1NoFP61jD401gygZXSbzVy3N4IW531fkqNxIo0BmM/kqWQo4keNYfKvrKwKz8jqOFMh81pLGdkoSls5g1o460bW7pyh5rSipsc5RM7lv+kv1u2gtbDPGCd7SsuT9bKRq4AsA+iGhXqD7uRgWdXyps/LAFFlJR2yjiyL5VARHJs8QLAw1gKursvkXQ97x1RQnQxLUZ2cGPQY25JVIB9esasvCY6yqO0w79Z11kNOM6PFe1ZGRDubfHRrFUqfSbL41KmqSgjl6pEwWvgh2BhxrsN+2BBUQFdtZejjIqAsCKNdawuVrKMgutRsElNK28kMaxuQC3+nZhPIFVNsaRZIozdoSL9JY7C42Vwv2275JJJNXdIU8za+nDUe3MYb48HtJf5vtmm+1U43aN/HTgJm1eX9ShQg0B0Jwau5AYhO1+4OoX/GMER3RvagC78v8bc16uLvWTFIFWinht0dF7LnlTGpA2RqVToDoU6MKQlUqSj458ErJwWHKQQbCNoD/r1pgZtIk/KJEhq5gWrqU4MP/jrOQLiRBWkHB03FQoujDwDPAc/880xBcnLvGLl3QA5gc1LvRFJ88ok+nL07F1OhkHTdQf5fEzdtoHXxDniHv/yWYbcEm6/QyE1JtRijnn/pp9R4fhX7CyfGpjqjsOZf5vSP2YKf9nw0mvXpz/l8NLcI/dl1Havr0p/esDvqjunPSX84mnn059gdDPszs8UJkt5o5HZpzsydeV320x7NiDs3AYaKSeXrDYEUT+ZsNMd/KLZD3BEZc85msxHnYezNxw5PnYwnQ/rTHcwsb8B+9ifupJ9zNh/Nhh5jx5t5szFjf0I8hwxyzvZy7lwSBG/AiwKWRlOWoTgt4JnM/cuSz+LBiPNTSJ4AfjbCUd2DRqZuS8vgLoyfwMjme4Uvwwi28gE9UwhunFFoQQu5gtAuDFya0fvAvWoMBl9j0KBpTvc0igRFQUccUW7YD/y3Oa2nObIUmqWOxemiEZzTZh9nuZ6eWS2j+Pe8JWXYNNUuZ3wSjPNmyY+Thd7BCAVMA1BlN3c4TbXLu9zsl6yXEqkMSppt/OSl85I5j02xo5ckLg58d20+DTmcFLqDAZUH/N0UKN+rEUHEtiBi+1ycfyubJW8Yu9SU9VhVZZ03FRO8cM2hDZ9BO4AjcyMJmrUgNhXIRRNYAH+8ZZTarixS3TT9qZqrdmb++JZjIrbv3jz3y9KWVd9y7Xs1nfWnaLpCorpGqOo+9f2xvqvaWzqd2kR84vLQ8eVzGB/yprK9k8GdYhRMsXvu20/v2tx3bfB7N/m9G722g23LtIX+Z98JszrzXLQ6i96bN+exYd25Oa0/X3NadXK37t2cdyaraU6x6QqzJWvgpty2uY9/d9tDW2CJhJSHf2QV0GhTxqmgUlLTbe6tL4qXtc2QNMylA85CyqPDZin89lEWpdyrsuwzOtaei2YnBwF+VsAGTK3E+2J9fnrrU50IK+ImOhnlYZsWjdu0hMDN3ZtgZ+WkvvcsqtRLdGJfOXGaPL55iqC5Y07lIjpcKL/ReTkFmq3bMnr6HBDsYCvEkEP08e8eKNBWyAGHpBj9+2D08O/Bx2HC30PUl8L7FGPMfniNfjwIdO5AK4mdlVnBXxu28ZUSjyyspxw7jVdkixaG5NJJ/WuSLyEeFyucGczCuQRze+VJSxCqzSbEa3DVN8fp4HhMh+JmJ1kGftowTcXUY0jZYuVxSbF4js5oQNCLJcJCmSremUT4XBzjUcrLmO6QuIjJkjhpchHNcbV6FQTSPK4J/Elkp8bBga/OebyHJylw0zIS3yOcB84lY1kICJZkCIAv6TYdLh2AbcrAueTo0F2ujMpTeVjRBNeYyBS2jN1EVTFf15XLg413KZWCKsVq7QM6Fpasg2IG4Y1TzCKl1hJ/V1oIijGU9admvdVSnkVLplLV5CeM62dK9c70tl/h5GkzNAbheUV1NQ0s1l878IhmkRjsYsBfLI1PaWlo94rsLIbdxVuxlIJmBFuKexR6z0LPd0nSUCPZPkvHHwnoCUVVNiBlA/oZHxLoIJytJQGSASOBDjvTLMg5M9GOMc+pJQOzDomdwBSVjZbRWa6SqwyBMmpqY09lkqrFnrGeb5rh1TzjqM7aT8zzc3niy4BPDA6VvPeXF3QcMpWlHoHdtw9uNeCbIzmZhB4kvtV747zg47uWy/B0VLfhGG1dWbgzCder6vnepT1Y46mLZqUFkTJ7FPEONTIbCyd5TzwjWqVN815sXgRR9H61LHGbrd8Y33xjfMVx/cswitFBpKNlTetU81WuDlPVZDVLUrDRaN/NVZDxdkHXqs1zxV3MOJVRqzzHEoer8H0YfQgl9iqsBgFPLKxqXtpF9tgFdaJH607fbTuIou+zp9u7Tk7zPj2grFd3ajZWpML6rq22tc3qe0hFe22qHBGxuOO/npqKUxEFpEPiOIob5i+MF3HsN/k8ot3ZxXcBsALeRX6Y+R7SBtNH0MhvlsQtrdJegPHn3PwSpn7wd7bhu+ERjP3R7cktw6FgmeyE+YttD/fTm59mCYmv6d4dvg2WxAmhaFlWowEmbuyTJN/HnDUXTz+zzju+gPgaSwf9s8qTD9+J/cKJ3+Ou+xNDYLfz+4rEN2/AnHXTKH4UBA1T3bktSp7VrSGuSmSeEAmwvyiFyRoEIJ2YLKJr0mjqNLkkoI7nJ1CHEC0PtSlbxu0m3ykMtUjSRyG4Dcjfd7GzIMIWcD3tiP0QGy+zY/TbumcrP/Ae8U3m3/mXq1hpeZeGSrI6b9MTmb+LXanLLIp6eUf+MstpHrEINkZ58Evc6kq31D92EjLsF0BCYgn2KYsTSaA0TYSkk9YLmJdUykqGiIPHDzBi4NE96BxeSBRh11nfFUCLNBHyRgN5o4VMApgAPA24kiHiyCG0AsVVtoEIJq6/hmHwFYlxD0iBICUrklyR76iBjYsPzwTrV5JqFdB0rxiksBPT9ir2WolNClPlnJ0xEKZJCsBGFPXEg6Al5S0UCmE/3EqWHUmoJJphplfomOCo+i2bHFZhslouozgFG8ij036zvFed6t0FWkpyoeyQiKKVTVFoueSiAD1cvg8eE9xVHMNYLScmZEk9GEt0YZQ4Ur5Rt1DZThZ5uFGTmsoGMyyATZq8/CLwn/ODrB5k3/m2fsq/lLhR6ynXOI1SJ3gSBYlS7yj4FU9J0Xra50UGqw5IteTDKQIo1Vu7iS4BAMApgIW4GQ9i4nZcyOmA6QQGFf1J7asMrOCIiaywS8S6Fb8fAskDKmIoC7f5h9IZirzmjFqOJylcLrnX0QdFcqC5PxD/8iotie5mV9Hd3EV0Nx8hupt60fHKFb+3iK6ouiA7RGxWqCKfclwMqr3BEfgnatpiobebavkog/UOQlIwBEmVCz+jcgrB+DgXDgDtqVy7Drhtibhu1/Aid7WAjtdxY+Kk5NuA4FfDZKBmfoaKfnboKUTc+F2o5oHRxbMR2dZDCfyKSjaHp+2hhxeE+oTi4iIpWacSr5wqX0KF3IbZ9SiLmkGCb7emB1N+5LuypSBxEbAH4xOnN2xC2TQSlkCymfN1sftdjDfXrVtpN+dpUTjPsu3KN/lDgSX4hXAWQNjZX9oazrf2q9vCS4vLbEMMPf7xaJboSswz5V0JOaKzrkGkmSWG9Q11wuRbXsgoNQPfzU73G6NAcL897txWhaNfvOcGJN39/iIjWXO2qfij47uFpbYKEbYKobQMq2MPmtNdqyPxhIkHeHrvsDiUtzUSzhfL8ISQRnX8UNOEd9EzdvJoIZ7S2VYpqzPYoUVqWriN1aEtjJyyry0RmD3N0Q4w3p+pxyeFmC4dgMGjVEZg6bBFTuLgRPQQxIH4Yf1APVU55C5EMZ8I5ghqIBubLGVDxU0Bf1MPz4exYuo38Ex1QGB0jvHE0q2weMKtAgabp1aCSz8FBwHwqeF8lsvrvKx0WZTxRLKG+QBg/JUu3hhHxldfFdkV9DTb2MX4gTK33GF7+15p0U7SUUkH9T+ZMeVdO6FLnkSrMBV1775GFbeIMt0C++ZAaNzcwMHUws6hYJlBJMNWarVksMnqW4wIIh+nJ4KhhmvFpUVY5VOWTbkacjUV2cm8t7fyDl1T4k6lhTIDKo0KoTXr8StWnZXPGdhb76vXjTfKcKvozkkuDu1mBKW0UlnivGHslbdsSOPAulpt7+NEUf+nRmVxiDo4KXyaSn2t1NZ1vbaiONeyrq7rdFX6WG/R03Wdllbq6Fqro+tqHUMhoYbqpdSsQ9aq587qIpmva1Up13VKuVdZQtVkLf+lGcI7gR+SX7lTYk9rAJM0jt6TipX8KspPnCUCJ7+vnJhshf6PiJpa5gJXd83POsXu1U5svLK6DSdsFT0/J9ZGk6ylGf22weS5D0UdtM+Nv/4VzVTcmFCDIYyrVShVU6puHrUr51H7yzz6ZR6tmUdPP908urfb5GlXTJ72l8nzTz15nn6CyZP+XwyGFdepkFSKUzRC8iH7LW9WEjJwQtcFOZp8HTtjoSrSVlC6Y5QM95tsj4yJkbvSGq0aEcwlwfYlvMJlLbzs4KbqjAP1k4XriGjwgQZgsokMr/iwp5rj9SLWTQWWLBDc5Ix31yA+/4mpp2JHLmL8xSVK5wC5Jw+3p9IUJUS3BSzNVThiLLPeCstOzSWrIBVi3p8sspItDjMiRfwfOuhO4RSD8yaHxCm96R3DNxLSRo6ioLZjPr9sqnSRgi/Hbwrp4zh/ypyldrtc8brFpRyGMmtsH+Z5UC/yVsEqkeHp1U85Dv0S8dT6FZ9fZ/SmpbOuUewpd3iJiIec79Ie/eISMiRwoFnzyToK5vP+wZBOi8UiebjUqLU4/n5+FYGeJ4oQPrUqcnN/FbnZSUW22quqjogItUpSruH9lERC/PMoSXbJmhofbRnaMGeLs1MELc9rFyiL+zAxni9cElRswMLQZeg9AYfTq1wG9PxrsylfpOiHlKqwLqeuqLCC706fUpbQ2IpftsBXkKa3d+arkaZwK6o51YLmK5E7wOK9wN/Ra4GpN59dDFwBncZOmOBO859i/5IFANJoCWPAvIo+TFyv4mhJ4vSmYfoL55K0Y4La74eXeL8L3XMDQvLM5g4E2u3s9tH2H1G0QAL2joh41Vib3oKcgGVCMfvLtZmJmzWHLOq3rhO4jWsnbijlotX84FZcJt4s19mRQJFS3hK7kWLgFbSyu2npKQKUGZgHZqYsDH/nRhLBJTnRS5JRODpWzfLisyp0fsHvcygrr/WD29LAv6G3N3YGZIHjLNQ3q3AFvZ+jpUBuvRO50i6GIHjuzEggbu7gq0o0/ZVy7oHPLbj+voB+QWFoSI0unQf4Ka2f05RMg/DmqzfYkVD6SyGeJkKxa5If06u2Ec7q2FAD4cJkLRbqMLv3SBiKTCoavLPNpsuyDaszxMG13IhN1DiWIXeJZnYWZfP59wzR+ihHqutHTZ0kxB6R3desFVkgqmN1J+Rei9D/1PbiLcUI2Zb1Nejbg1sf9Y+qnx6PjydCbbdxopquRYRSwxxGQZ+wQT5zImiGjhWEfYTXr6PMhBHB0Co61/Lcoy1NWXJ+3sPY7iIxc/PZ91L9N6uUNGrGUQrkcdhsTyyPXJpa2ppxmWtUjMO+tphYmj921jidpraN+2v+nTW4bFjXqrBGOWWZ7KCde7IFDar6/N4qUaBv7V0C6A4qVIKeFaN/NVCmBXUwMPZna0Mmtnc1pKS6uVK2I6qVqMSSViqqLPb+gvJ9hwgNBWEyZjC3FOhIqEcLj1lhEqW9aVbtAKSBl8fRmiTVFvxH+QhFAR03cJLkuZ+kHbBYwNIN5xGKkr/BkFtOdwkPcTB2xWHtNlHuBbFdxW84U4riKysJAuu7VBvfJDGFK5yVsjiNxtv/Ch/cFn1kc/ZW2cCDr1+UWKsqlD6bIfRKhsyjqIbJXtIwldxsK05ZShyA7RGrznbWsrhtBSBJyVJc29CLgomToWyTmnlumIqcmN7cU04MOZcTfRTFVDKrxcQBqsSUZVeKiQOoYsqTyx4W2qt9y73Cq1USPkEUM0SdfBnRsuxcDJxCzhMYb/CVEvpWBblOJYO9JAsA6EBVL0nKkorQh6Jb1YB78kp19bZlfr+ESknq3kJcqCL+rzY5jD3f4q0NOBARGC4z3WjphKL2q/sglwddjOIEQOETDrmcYjH+ZCkdupHOzO2PBTsOZ5YOruM93zDU4il+vNGmM6AH6m28gUZcIZBr8NOST+MS/3nZu7AeURLiwCkTzzUqvzRcPC+D5lR5INhxGCiQtQOBkM16umlrcmgnN7uaLN69TSe8KWdmtdIcqBJJkPRRmsb+DCbzhklbuCU2rRKIm0f4zNanC/BxitqZWwGpd+ozqGw0/x8RvpzVVonw7EIA+nzzgPrz9FjyPGqW4xj5G02VHaxshO0kkIwwk8hLGKiwnlmqOVV5UCbyT8cDDkZlHvAyLAXobuG6svuhENo5WlemxC2w/AWnTxiHzmhWmZg5gByeE+5YX7+gj7rBxKGMdBgAwo1TNNaP7QXz0GN8z8sPL58EPjDzWjoTLCz5hOAWvc7E9eA2o8RdlXZOmsZeQFBvtStCWQxAsPkrdoLkJm4WphNw5CUeDtKhPoLiQOVca+DFWKveRUJpi8e0sRbUzuUyVEhtESZbwsnxM+07LUTHUpSbt7axq9j5MvcsUNGgQYqmQbMw+N5oW50B2l92p0sWyg0RH1NDQ9YTHr8tKUshBK4teHDKvRIv2dPsRbk3aUsOnKi7SiQ9n6+CrPIZ7TpNLvzbbXrMieYIuJ1nqzwFbuRTU/P0qHi2I4OiVW5lhdEv6TWONFrqsCC5QIIPCYfK9ag4PZtj0YwCj35KmNmTkmVUllPgsm8BeaN5nKVu2UHQAxr6bRuSSMThqG7BoRTuYq9h0ghcXgAurQj04VOKb1XQF4ZOWYLi0Bnz1aBtxPIVBJEaD+II5FhKQW+jnErHiQmooflCtwwlxb4WprWrJUwzBPO/i6PFG26Z6o8RSrsD8KoT79WaXZ7B18NpauODH3rRh45HrsHDoKVSIOlm/goYfOpK2tyjFmMbh2JR+LmtOPnxn2UO8Mh7t0rSBYsBKgVVUq1+bsdNklfrn6nbR18oiRN2nr+h2tCSHEQspe4yQVutDmJL6A+VSjSNr4X3NE5K15t83JKu/v7LO9F0VmlkVjQTXbNEPRRrKG5Ce6htyapOtWXBGfoYLVC4tUq0sKQlNMbfByd1rwS9VXsM35RCPF++ZZir1QLRX2Bu420jJjSgTJ8ofXBboXgbb4lDkBIWCbhbXyrfYGXzO2OUEIB7RZ3GVo4uzEa68UCcrZSKC3GL7AcruBx4qCp1U3Iiq4MWciVrBrf89rjAY+vMAnl2lRMkZ6D1dd6TWi5x4ygIHks22S2ddHUl4FuGlIOWMSNXzrWPr7iaeL+SE6bmRi1DdKRpOc/CNPq7Tz40bjXoQDOI3PeQEhIHlKgguJEPvZfluYhWCdUMlKkaQ8uuPsJ9piA5ccMphrZY5f6zZRQfv0ljWoapOQGdu6OsW/KXemngNwqJcH2k4sVVAVbdBJ2fhry4Vu5IzDcf5TXbpVAq6Gr22Fq0ZKVvcTfppll6ILpiDa1URqpYL9uLoAc+theBzSEviOmvnZTKrhZwzZU5zekOiqCKWghOgGkEEEM1Dy0m4mULZTllppos+SAnoEf+IYsSKNg8XYdOLxQrugP0dJGPU8MyvvlGEpkIfKAAM9dL4Fh2EBVp8a0PpnTZrQLDvWPFc9xt6ViuVlZ1RXc0zpuWTS0HOxV9UFV0VWUliWwqWuq3vKV+yJzzmqb6LW+qHFpsqx803rzCG+239W2VrRd/ZGP99jGNVR5e7tBWv92hrYrFcbPyCMhO01fELAJ19tpxitlpgtnKCI3TbZtDX2dHDe45k2b45em0fCoGAe4/OV4UsyIrdKpep3pxXb7+tZ4L/c0eeQBUnnnMJ5hBvCMMlivauOUyZLUypetOdH92m6wM+bks/ZEi+UTVNplUHMQqoamPAuhvCNlJVz9cERLodDUbH50ghXLLgS2PBKnzG7oYDLfZYSkFV4jN75Z/SuYOKI4aEi0eoi4Z7U18eKHIZ7SVR29qLX1RJmVA3NeusMtdA5g1EO458B3s6LLn99OJiBKhqXoQuszRqdHfVrv2idFX9U8q9biCXbzNpnSdjcyu9Hkgn2DUHfWrrMex0d5akYNtFTmtqogf3qki7e0V0SxmiiR294HLP+639FJ2Km/VSC7tmC7F/U2M+2EI1zhUdPihMuY1FAlJ0FIclwWbhcL+U13u+ZSl6XzZor00t9UunSTxr8kRe71lI62J6XZR1DWgNoBRurCWBap+CslT/mbPp7uuVr5dSLxoVLP5lQUvOFCWooHLw8UiKEucytczla/Q2/0SvdI1evIleOVr80r5jIVU3S1WfTFe+V4k+poEe7SBT/541wRoCW6JOBJeWQLN4bOucjI2P1h0j12klUd4bNdZmlVQxemdWrAFjTiyNjetzmhAFpWwgsXqh3h7R1v2nEvrCvqjJnpgNI7nAT0MbLLrn6th608baSRWeY5ChcRa5Z652e0MzJ3iu9vC11V1XqVYIG2k5Vo5srFXdzYsP+QldYJa/RBBqpRDhKnYgLynLDr/jNNC6RJzzPk2KPniO2p8Va0zsqJxxJPUunfFvcslMO3hFwOD8tgFwNXMlzNtus8LZqKsstlNAxuyeFtJv3bzdgk6C660B3U8F16sbdfBSdvfXRKKF+uUi5b03dYByu6SOiFsVH1gM8CfRyVYff9tWsFjPf92lcgmflkrqrYs5W1ZvsOfn/G95W8FKDbRUTmJ2XGa6/SPdIkIvZnubZrUCkuvfGoCvI6i9GXkkUazcxUlKXib8zDpZAGo7EZG+Dk9PgTL2V+mp8eHaUxI4sIU0I5XYfuKxOT0GDe4G3SH1cn+PAo8fMrjIgTK+6fHVE6nx3RhyUDL4WTfvSLue6jDvhbnIo0uLwNEPaRIMnl628eFM5uBF7xfl/VNkE7fOesOFYTBr4Vv9FvoE/7rn/+zbXWssfF//4/VGRpn//rn/4bvCWT+65//y+7Y50YYhX+QODqyh1B1KOM0/4vXRywYHP7wgqwhxSOewhSe6AaV8S5cVAFoxtp8yMR9mjwvE/9FrjgKsvIKxv7p62wuZk3X6XQyvrXNRtuUNwooauC7VGkPIzclaTsBHGexf5q/bMUUn2oP+0LlmebBizm90L9WuaYMrAM8vImiMHv/gpNIlzw8gOOi9OhHw0zJYomGBdIhcQwCAxMOj4YUxiZ7v+A/3vz0skNd3QYS7PAT3Co9VnezKfVqoIYo+fsfhtwjOrrXKxi7xdMX2Ku3dJesLlkzSlXar2mXd0kE+nO7j7sX94+M/acsQJZfgQP2PYb1wKM1HLpU3+wYP6CJecjScdMVe7eAhgT3W8a+8FoBUvz+x8nL65F9E8d9uxs+Ofzt5bfvX/0Yk6ufh78cPBsfDn7sP3/VvXlzPRrEi/6Lw+WjD6vr/nL5/fPRu791308W6bV/HYyd70cH4d9e9g8eLZ+/ur48OcmLog8qYEm8R2KG8ngB5j7627ePy/8xOsLtCQB6u+/Y8JeNBVjwow/eyH5xehZSzm738eAAkgVQyKVdmH4ClGEDDqbSUwOQCu7OPnQjSuocft9U0bLKtKwttJTjSkgQk+UrjGilMjEVlx0hXZiIIbENQ1S3Pxr1Bla315+MB70B5jlrlIMuiwdSIZt5cPv5tdPIY8eyLWswGY77A7s/mUx6Ew5Bb6Ouopk9y4x1OJsM2YGAHn0v1u712Wffpp/DEeQM8T/6OYafI3z5ln1OxgBs4RN97FldC7+RB9tmr+jamDDCdyAHLKELCV2rBwkT9l5tH48mYAndHoOwung4gf6PFTJE9Ak+G2hZNMHGIwwji6ZiwmiALGNJY4rRw5/dXv6NxHCC4OBY4hArPaQv4vaAWh8p2vb5OSqgtHkfRDRial4VA2ealbU5f/RvP8Oh6oWdZx9jNLVgZzIIvU+E9j1QuUJpi2e8UDOzhExRt5UBv7cUYn98Iefb6/pfoQHjvQzHMjfnG3EI5sMt/ITJ9VSd0NlfO9oYylT/D+Mp3T50ZDx59YthcWLfXKZTfSmMAZziqdUJQz8LLhyhwUH2Dd9DYyi+yGeEMZnZtjPrj5yx1e+Pe443sfoDazCezPrEHs45e5/JBDMoGPHuaYq50QKdmos6+4jmAXjgLBPMRLNN1yYKjb8YHOiTmHi1DFU35mc2BWm1GmdnwPfY6vUH3QmMIvDR61sje2L36Mekb427k6HRMoSBvHhi9QxSu0O7P+z3YQArRnIGP+gPhr1Rj5Md9e1h3zBEZLtjdcf22EJ4ABmOJmPQQ/YBo+W4b9m85MFw0p8MhMddke1+v2eN7eGYIYPWjmycFgC+D5UYjwxGqW8N+rY1xrfSuQNFlZQ3aLNC+h/XsQa9ruNZ7sgbeXZ/NBjM7NHEHRH8Hg67Q+/zdqzP79v0bKr5g3Gm+uPBF+/mi3fz5/ZuXv6cHFyP3/3w6/PDyfrHRwfjH3/8+Tq+SQ/j/uT3P348HKWPnl27T9dXfxw+fWH/fjhY239/FYJHFA/D0c/9w9c3109eJb/+8PLd6PDdVXKYRt9fv/ri3Xxy78Yadq3hAKaqAdjp9ngseDearG3eTQ/mmm5vBLPPcGL1BuicKO6NhugX9+aLe/PFvfnc7k23O3RmExvUuNvrD+f2xLFHs9l4NOna4OiMvC/uzaex8v6/d3BG4ClMrAnzB8Z9dDT6zLPo2X3L6vKc7njcG3QHhuJmTMB1mXT7zM2YgPYMuI/Cx/4xdX34KK8gjwcwqnZ7rDAYnUEZx2PuEIHrMjBYzqg/GNu9sSG7Vj0oqzfs9yiIPRp3c0pDazQcWCNeh353NB4C2x/r4Hy0ncuuS3nm0bsaC9Ns7no2yGoEbmKvP3aHznzYnY9G3R4Z2E6/P6SGyc4xCRF4q58lAm8dDs6n3ByP3BgscK1BntnMKCsOxMEVK3m/k4vzArP3S6hPgojuU8i+Oy4m0GLosyvTIicmMNC55Fc/vWpI6DnReexc8kOQyuLvU/75HYfAOmTQ0gJiRhjzNeen88ZVzk+zB12zA445Zfjky86Pb54B8QwbXzaYsqeAo1Xskqf05pAKGf4Fx/N948BQ0Pm9PKJYcmqduR8nWdm0ZoBQ5BYOykZuCa2Qt/kj/w+O9cyg</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5ecb69e45849446495cac4e3367e5c4f\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
     "\n",
-    "Once you have the state, use `nnx.merge` to obtain your NNX model and use it as usual."
+    "\n",
+    "## Restore checkpoints\n",
+    "\n",
+    "Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.VariableState`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.VariableState) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.\n",
+    "\n",
+    "At checkpoint restoration time, you need to have these classes ready in your runtime, and instruct the checkpointing library (Orbax) to restore your pytree back to that structure. This can be achieved as follows:\n",
+    "- First, create an abstract Flax NNX model (without allocating any memory for arrays), and show its abstract variable state to the checkpointing library.\n",
+    "- Once you have the state, use [`nnx.merge`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge) to obtain your Flax NNX model, and use it as usual."
    ]
   },
   {
@@ -185,7 +203,7 @@
     }
    ],
    "source": [
-    "# Restore the checkpoint back to its `nnx.State` structure - need an abstract reference\n",
+    "# Restore the checkpoint back to its `nnx.State` structure - need an abstract reference.\n",
     "abstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\n",
     "graphdef, abstract_state = nnx.split(abstract_model)\n",
     "print('The abstract NNX state (all leaves are abstract arrays):')\n",
@@ -196,7 +214,7 @@
     "print('NNX State restored: ')\n",
     "nnx.display(state_restored)\n",
     "\n",
-    "# The model is good to use as ever!\n",
+    "# The model is now good to use!\n",
     "model = nnx.merge(graphdef, state_restored)\n",
     "assert model(x).shape == (3, 4)"
    ]
@@ -205,9 +223,31 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "    The abstract NNX state (all leaves are abstract arrays):\n",
+    "\n",
+    "\n",
+    "\n",
+    "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_5eb8fb6116ef450e8276128216cc9262\" ></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5eb8fb6116ef450e8276128216cc9262\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtGgtX2krzr2zpOVf4lPBQUPFxvoC8bNUqtlp7e/iWZBNWQhI3GxDv8b9/sxveRIrVau2FnlNkM7Pznp2Z7K7HexbZVzgjxNMcl9SZ43D0D3Idj3Lq2DnEiIU57ZAdZDg2jxu4Ta1eDrUd2/FcrMF6t0k5icsfOeQyWLGox+Ny6zjvubBqOzYsN7DWMpnj23pccyyH5QLUHdT/1bAAAPajOm/mkEE5gNmc2HwHuVjXqW3GLWLwHEprTUHEJvEmoWYTVlJKRmxjc0yB5yFa/494h3q0QS3KgXPsc2cIG6c2Z9T2qBb36B0JnvbZvd9NBOrZHaonznwbaDJY8zRGXY6EfHsr2HUtqmGhsYSjcSKkZwS3V/aj0djePigU6Hkc6cSwPbSHeJN6ikn4GWj72NFJNKY0HY8r8jmIRjiqu8QWIqua2FUgffse9qSCbd0i8Nj2LWsnoKAAmzXHsWE12nVYK4bGeXAuYEk8mljmVBOLLmGGw9rY1ohiO91oTNoXCERnnqB4gLSL1tMx2IcaKDrFtWIR2+RNtLeHkgJkLuuMcJ/ZoHdELI+MGGv6tuBsemuvSQ0u+JMA4o97+PcAhSh4la07XYWRG594XLVpW5qrxHCbRAOdxMQeOzOEXN9rBmrcCZFxQGIvEGOOlIvzILgIDMkd07SCqKzLyAFvdcVeYoVYfA2RDjh435KCO/lbaZGeUHqERQRDfWBFs7DnfYTg7O8bjQz3rLfBDSMD4vcx0Ce4v/Tx/d1EWADotIPkhnuRyfQRQRw3QFJyuxdJRiB0GZ8FcWxgEZRhw6N5wRCugajAGcgegWAM0pjMI3XcaDDSkf4j08r77FYaJ5MgVR9Ac9ptQByDwPIjhJ8CwTnb4dFc0+kQFguBnwSvd5vErpNbF0xOdImqGI6l4wZIYINouSb2ovsWbhBrf/JJPZAzIKc1idYieiyG/hNDD1MFDYJb6mMQSfgYxgjC9tsNwsYBtrdSmewIwBPZzxynsZ7KpDICYIK9B44E0FeYFACtU8+1cG+Q+qcB0T6SWsjlGgSyChnjQJOfnVB6QfqPp0T+7x8TYNchLWrLQ6FhOeIceZCmtOYsZR2zlkewCZ5qz2I/kzWHPAjUcKQB/ASH8hzLoZW/05mGtvKa7E0iPchk9gWYFHYUhH3mCQO6DpzmhIXQpd7zkZWhIAnFZf7xHvLx56E6Eo+TWz5LRaFe3aDM43XHrgv3DwmteaGkpDMimkJNhZ7MfmDxaRaFVG3MTKi/AjZkQN8/kRqkNLfX8DmHwicsAY0ehzltBEWmoECRUNSGA/9NUht6ZKomXjnC4BUUW6jWazccy0MnPhfy6qgQYMK324PAiHdJowX1rUT32nC2NSENQ/Vpc0Cn2CP6sCp+T5Li386smwfYsmxNKtukPS1lEB8hUoSnuxGm0sVeXYOKFhQ7xMcGnzhKBnl6Hs0pnEmS46pHHcyi8biOOY5jGwwrC6PY+LIgIqo9hu2BN8ttUcpDBDQGpXzc8fnjRBlyAIahRH83yYkkid7Rtuswju2ZvRvMacGBL1ZGyejH2h1DG9PnwMz3iqidgDG9rkHlrTNi91md7G1gz0nAZ6s3hqHTP0gnQlXDlhaFBgyq/5R7KwtGxeNY4A/5/WWcNKDcgbwbcKI7HGQXXIwr78bHlg31dB06UYPewiYTYbIlwwSaB8yA4S5mNgRefZDYB7YwDKyl1kMAXai+/xl2o6zffIrs1VdSfymeVGRaHbW8OdmwYhY3GdYpmC2aWs/oxFxDDni0SVASuMtqzbXAw6FqFvlCLqG+lmdYmUmsz6b4CZMPhLlXZH0LSrKwCynqx2Xe43P4wxSC4k7SmCizH4B5Dj7CSPQFnQSoQyk9as/C1PIw+PPYbA1NDXGUydYOLcjwtJaHpqgDHrDAn4z2ovIuyP6EUecbchGqC6loEV0uQuxHQj3HFEsMFtA7lTHcUwzmtKO6o/mi41VEcvWUDrZ84kVjMcVz2iQqU64YSYhvJSiPxDhiwQIpsoJiKDYcAHlNQriYEpEuKtRqNSFNTayJmY98qDACImuk1rO16P/+2y/KNDJI/o8v0IIzQ1ASemRtbPXXuv2B44YYKHhMyyGfWVFRLeTE80TXMYz0TgPqkezGmp7cLh+Zal6Vn+qpqjryr/xZF/6vlFS1qM775NuqaracD3q1mC90v6rq+dfCoXpUzRfUknlbrXxsci9/RIm5Xjq4TH+sZr92aq5PPx1lzlOHl9WzL0edi6M7/qlXKhVWL8zWOc0fJJv04NQ/LOrl62SlkTA6Vd29+ZBt3lxQeuof2eVmxfjM1c/Z/DHbUEtVu1XMap993149y9xoXqvbMUpW4ubWLDpbZuOwW95KVdSErZ5lPjJ2mDpbNe+SZ3pSPTRS5vFmoVu+TptJp+efbW62i6lst3K5fWKaLjlv9TZItXGX0RrspMyxap5Wj7sH2Ot5p361enlRLHXVT6du9av+OZFYNTfPNy/XedL48OlG7WRgz4/q8aZ61FXb5t1ZbdW/qpHi5W3ayGp3xxtnlV7Gz6sf7vLXbsldp5XTQjF55X/aqG3aRv5jsVI6aqt0datTTDftVHNztfGle3ndrbDOQflzwb42ikWTr55oV5a1mdkuHHbzW83tjaOjcm29fKWa7WrmOn+6zc/LpLJdzOer5fUDc+Ms8VXrNdQy2PTLh4R6WsYqOSpYauWueGJecTOb/2SenFQP8i16miGl/GUhX9Jo0m0yx7XBN9yr4kHqLtWqGQWDN3sf7IqOS17FSB63y8XjbF5Xb758cTH3aldtXcd0O23cbW98ptc3WbfNsifO10KNsnK7c1her13U1kvFtJY/Nc5XK5bjljdKXjeDzZvsFr0itWPLvbDzlSrRjxjxL27KhXbqosRatdptJp29uPC6KnAUQ3Liy6Mr0q1XRB3yP/hvGP1Yd1yowUYhKefUiqLMgVgLYvY77DV/8teUg1NZJgcVPOwN7mFrKBoU0pNjbQjBc0eEL4D1C22x5kF6EFuIzkGU27iLKUc27lATc4cpsLPbcDDTlS6jnJxDkx0d7QXC9vcazU6h+ItGxtoKMTUFKue0TaD/iA7G6jN4jLShR5hBvV9D6WQyKWs8SL5Q7kVlgxxOd6x3iIyYE6OBQQYTg+YIeo9KmFqQ2LiDBPA7mdmgrrOhPoZsTEFnBOui9Vkd111/AvyD2a9oowbD38kZ23RFGtnfDc71XWq7fv+kiciTvOHcRkI36R/68DA48IEJiTxJd/Kkjez/ZXHBNkDMh5t4ONUtRPYNC98qtn0rehpOLNpQBnsGXzWxHP1nN9GXanyzmb4tMu/51MPxeSwI8/42vbkTdAyp4EefgRx6WPvPquxQGVdGbewKcuyCcOO9lUfGsZzex1bQsMfei3xb6cu68j2C5Lm7FxlrwnPorxvf4TtjYMHCDpoZGsAxLz0ZMmUT/u6reX/SjL/Wbg2KvRcx2qtERgdDCQWceFOh8aW/HoTIK3nPtxWh/IW9aAD+it4k7f2S7jHnlVUkxEFGz96jB70otK98skd9wgy3h0shGXce888Q1cGbxCCA5AMURHUI45LV4Olf5pxwC2f2BQNDEQ70yOiQOM8YIhMeGNlHUK0cgEuIPgxzOSdD2cwWcgyU+OwR5iVop3cH2jMTyW4yIbSfaPhgtES9TnTKpd/XFbEeTyrbSjLu9tbjov0FLnsBPFgrMbKW25vibs7X2qwRh1MF4H1p6flKnBdycmywF+Ia/Tfbkf1jMOICGl77ja0ghXysGSTSs9oh+IrNLxTGMtcbd/6XPf8Xd/kf1o4twmxivU7Jv6wih24RmOERfjRAWFaSy0ryT68kB77+uApjAmtZTb6ZA/X1rf2kivKFhy3zE9U1hPvUQLGJXXIg+Khx5mv89Y+8R1ask2ivePx5QpO/2va/i3kUKe1PGilA/r0mqMFV5cj+xh+QKMeU/C35/UlGAvzX7E2mrfLmmvBQnaaeaJPU91/Smv9Rjv+aqWlxl9dne6bpA912Bye2BI4OCncH8/V08CP21sNC/7nSbgz5FwbEG1fu6xRSzzH8e7OKf+mhTdjX/VtN5y/z/v0ZJ6YBS+l/0S2J9GJWSi9vSSxvSYR6z2PekqSXtySWs+1/xWw7/RPvztPLWxJvsMz5DSy9vCUxodGFW6X08pbEWz3/l7ck/rQq8jENd3p5S2JZSf6bKslHvzdPL29JvN0D9fWtvbwl8aJ2flzFurwl8RrmWfxVZHp5S+IVLbTgLYn08pbEC9ok9USbLG9J/N6paXlL4pEW03+utFvekvhtC6l/2y2JVx3azLsl8caU+DJv30MV9pzuOYTUaWf//8wFOMg=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_5eb8fb6116ef450e8276128216cc9262\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
+    "    NNX State restored: \n",
+    "\n",
+    "\n",
+    "    /Users/ivyzheng/envs/flax-head/lib/python3.11/site-packages/orbax/checkpoint/type_handlers.py:1439: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\n",
+    "      warnings.warn(\n",
+    "\n",
+    "\n",
+    "\n",
+    "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_16a3675ef842438cb85db5d1411974db\" ></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_16a3675ef842438cb85db5d1411974db\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWgtT27oS/itqOnNILsRxnBcJj7lOyIsWKIQWyjlncmVbdgSObWwlIZzhv9+VnIQ8TAothXIamCEgrbSrfenbFdsBG9pkV2I+IYHueqTtuy5D/yDPDSijrlNCPrExo32yhUzXYUkTd6k9LKGu67iBh3UYH3QoI0nxRwl5PozYNGBJsXWSDT0YdVwHhjWsX1m+23OMpO7arl8Kl26h0V+aDQSwHzVYp4RMyoDMYcRhW8jDhkEdK2kTk5WQonc4E4ckO4RaHRhJSzm+jcMwBZkny0a/JPs0oBq1KQPJcY+5E9okdZhPnYDqyYDeknB2JO7ddipUz/ZEPUm/5wBPH8YC3aceQ/x8O2vY82yqY66xlKszwk/vE9xd243HEzu7oFDgFzBkENMJ0A5iHRpIFmEnoO1D1yDxhNRxAyaJeTgaYajtEYcfWdX5rnzRn39HzTSwY9gEpp2ebW+FHCQQs+W6DozGB65/lUDTMrhnMMSnZoYZ1fmgR3zT9bvY0YnkuIN4QtgXGMQXZlAyXLSNMkoC9qEmis9JLdnEsVgH7ewgmZMsFd0nrOc7oHdE7IDcC9bpOVyy+a2DDjUZl08Q8F/u4PsBDnHwKsdwB5JPrnskYKpDu8JcNR93STzUSYLvsbXAyOsFnVCNWxFnHLPYCY+x5JSPl4FLERqSuZZlh1HZFpED3urxvfgIsdkGIn1w8JEluXTib+mKDLnSY36MCzQilnQbB8FHCM7RvvHYZM92F9wwNmZ+lwB9gvsLH9/dTkUFgEH7SGy4E5tNHzHEsAYnJTc7MTkGoeuzRRLXARFBGQ5MLQuGaA3E+Zrx2WMQjGEaE3mkjTXNJ33hPyKtvM9vKliW4VQjAt3tdmHhFAUWX/zwcyS45LgsXuq4feInIuhnyduDDnHa5MYDkxNDLJVM1zawBidw4GilDg7iuzbWiL07O9MOzxmy0ztEvyJGIoH+k0APcwUNglsaUxQyfJnmPUXAk5s1vUUmnUvnOAFol/g+MdoeZF3SAVmIP01Y4N+jpC/UW0KUYchyfPGM6A9cF6DLqBMCtUEDYDocXwvzhGgXCQ2VShqBjEOmpNLF11Ykv/BqSKb53TC6QsDmE17UEReGZrv8jnmQp7D0ImcD+1cBwRZ4sbO4+pksPZGBL41eNKafkVDccSW09peS0/S11xRvdtGDQuZfQEhuR8645wfcgJ4LNz3xI/jS4PnYilAQjJIiNwUP+fjzcL0/HiM3bJGLRIO2Sf2AtV2nzd0/IrSWhZKk5Hg0RZoK/bD4ocXnReSn6mLfAmwWiiEC+u4HuUE+9IZajzEARVEJ6H46ymljKDZHBYoEwBtN/BdJZ43YHF5eO8DgFRTbqDXsaq4doKMe4+c1UCVcCZ/eEAIjOSDaFWDfMPN24d7rQA4HZOowWE5xQIwJYn5PZP69tejm4WoBaWWpSLrzpwzjI+IU0enufqU0wEFbh3sAFDtZj002c3uM8/QynnNrZllOqx71sR9PJg3McBI7YFgBmhLTw5wJR4I+dsbeLLZF6QAR0BjA/KTbY087ykQCMAwlxrtZSQRL9I52Pddn2FnYW/PdKwADfOQ+GX1bu1PLpvQ5NvOdxHEVCGa0dUDlhk+ckaizdQ/sOUv4bFhkEjqji3QmVHVs63EozqAySHs3AkxKAcN8/UTenyaJBlAI8m4oieEyODuXYlp51z1sO4C121ClmvQGNpkJk00RJlBYYI6LBth3IPDa48Q+toVpYj2diSD0AJn/M6lU/VFhyrPXSEmjoaQsibR6Xw6XRDGL/aTlY4OC2eLpTM4g1gZywaMtgmSQLq93NkIPB0TN84UYQiMtL4iykFifTfEzJh8f5k4S2BeUZGMPUtS3Yd7Tc/jDHEJwJ3jMQPAHaJ5DjigWo4POErQBh9+XblFqeZj8eWy2geYaPNJs2YceKfC8liemaMM6EIH98LIXPe8jxZ8x6nJDPobro1T0GF0+htm3DvUcHS7edEDvVN/HQ8n03W7ccPUer4YlnlwDqY/tHgniiYQUuF0SFymXtyv4pxTCI96qeCRAiq2hBEpMmkNBhxDGO0hkgCqtVoufpsXHeD9ITEo+ERVua+jo8f/9dwTKdDJO/k8HaNMlscM7YvZobDBqRmZ5syHw9RLq+Xaco4USn08NXNNUtjTAI/nshiEX6weWWlbFV/NYVV3xW/lkAD8bNVWtqsu+yl1Vta7cD0azWq4Mvqrq6dfKvnrQLFfUmnXTbHzssKB8QImVqe2dKx+b+a/9ltejnw5yp+n98+bJl4P+2cEt+zSs1SrrZ9bVKS3vyR26d9zbrxr1S7mhpcx+0/CuP+Q712eUHvcOnHqnYX5m6ud8+dDPqrWmc1XN6597PWf9JHetB1eDvlmzU9c3VtXdtLT9QX0z3VBTjnqS++j7++mTdetWPjFkdd9MW4eFyqB+qViyO+ydFArdajo/aJwXjyzLI6dXwyxparc5XfOP6gyr1nHzcLCHg2Fw3Gs2z8+qtYH66dhrfjU+p1LrVuG0cJ5hsvnh07Xaz8GeH9XDgnowULvW7UlrvXfRItXzG8XM67eH2ZPGMNcrqx9uy5dezcvQxnGlKl/0PmVbBccsf6w2agddla5v9qtKx0l3Cuval8H55aDh9/fqnyvOpVmtWmz9SL+w7UKuWNkflDc7xezBQb2VqV+oVreZuywfF9lpnTSK1XK5Wc/sWdmT1Fd9qKl1sOmXDyn1uI5VclCx1cZt9ci6YFa+/Mk6Omrula/ocY7UyueVck2nstfxXc8B3/Auqnvp2/RVy6yYrDP84DQMXAsapnzYrVcP82VDvf7yxcMsaF10DQPTomLeFrOf6eV13uv6+SP3a6VF/Xq3v1/PtM5amVpV0cvH5ul6w3a9erYWDHLYus5v0gvSOrS9M6fcaBLjwCe9s+t6pZs+q/lXrdZNTsmfnQUDFSRKINENZvE14dZrHIf8D35Moh8brgcY7D4kRQ9bkqQlFBthzP4Ney3vCnZEU1XA5BDBw97gHo6O4iGQnm15Qwieujx8gWwEtPlYAOmBb8ErBw638QBThhzcpxZmri/Bzp7mYt+QBj5l5BSK7Pj9XnDY0V73fVUAf/HYVFnBO6rA5ZR2CdQf8XHLfWGdT7pQIywsvdtAiizLAuNB8gW4FxcFcjTfqdohdi8cbw2MMxhvQsfQe1TD1IbExlzEid+JzAa4zgF8DNmYgs4INnjpsz6tu1F3+Bt9YV5GjRvDsz22eUQa290O7/Vt6ni90U0TEze55t7EIjcZXfowGV74IIRYPMt39qaN7f5hMy42UCynm5mcqxZiu6aNbyTHueE1DSM21aTxnuFHiw/H/9lOjU41vdlC3RZbNj83Od3MhcO8v1EKW2HFkA7/GAlQQg9r/1mVHXnGtfsydg25ToW78c7aE+NYdPYTa2hSY+/E/lwbnXXt7xgS9+5ObKoIL6E/rnsu25oiCwe20ELTAK554cmQKTvw+0jNu7Nm/Ll20ygOXsRorxIZfQwQCiQJ5kLjy2g8DJFX8p4/17jyH+1FY/JX9CZh75d0jyXPWbEIB7mfe48e9KLIuvKHPeoT9nF3MhSRcZcJ/wxRHb4yhgEkJlAY1RGCC1HD2T+sJeEWLewLBobEHeiJ0SHWPGOIzHhgbBcBWtkDl+B1GGaiT4byuU3kmij1OSB+kKL94S1oz0rJAznFtZ/SemC0VLtNDMqE37clPp6UpaIkJ71hJsnLX5ByGNKDtVL31vKGc9It+dhYNOKkqwCyryy9XInLQk60DXYiXGP06h3bPQQjPkLDG7+wFcQhn2oGsehZ7RB+JJYDhanM9cad/2Xv/8e7/Dex4xXxHWKv0OMrocdQ/U/wn/GCFYJcIch/O4Ic+/rTkMXMqhWKfDMX6etb+4eQZPT9RI2d2OTfIjPEMApFA/N/Hc0auoJ1nC8YWlExc5uFHM7OhXTUv1PGdi8hksXLGDJtF7OMEs9uoGzi22dYfL2LFLJgKppB0nIxJ5NsDmeKuiwr+fRmXjNyaTOHY7s/O/VH5c1oXfA094A+FrHlt1Tz/fB+KTZ52dh5GvqfXfZzc+UXGsB1Sm/FjqhDDYgoRB10/5bMH5m/y39/M0U/Y7315orc18LLUR93b/WufZknj2csUkORlN/oYUp5nJWU1cPUqrUQ6T1PaUwpq4epVVvht2grKN/xXKGsHqbeIMz5BSy9epia0eijKyZl9TD1Vu//1cPUvwU9PqXQVlYPUysE+TshyCc/VSirh6m3e5G+vrV/9sNUflMvZIpYkfWinMVmRlM2TSWjb6Y1Q1Z0Tf8lHqZMnSgkk9ELORlnjdxmMVc0ND2XLmh5YuaKyuph6td6mFK+771EWT1MvVVF/6YPU6+Kl5c9TL0xJb7Mg0ekwp7TPSeUBu3v/h+335xC</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_16a3675ef842438cb85db5d1411974db\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
+    "\n",
+    "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet62ziy4H8/BduTbkqxJJO6y7I9Xy59yekknUm6p0+vjz+HIiGbCUWqScqR26P/O++x+wD7Cvso8yRbBYAkAIKU7CRn9vs66UtEoKpQKBSAqsLt2POvjSS9CcjJvucny8C5OTLCKCT7hu+d7M+j+MIjcxLHxLsYzbszj9jWZGCR/sDpTVzL6g7t8XDmDez5wNk/PU6WTgj/R3qnHSeOnZtr/48LNwpTxw9JbNwaH678lLQBziVYULxwgqmxMXTAHT+cR4Ayh5T23Fn4AfC2iMKIYk8NNwqi+Mj4i0P/TI2FE1/6YXsWpWm0ODKsTndAFlO5xGVM6ovzw+UqPUtvliCQ2Akvyf45sHBN4tR3naDtBP5lCFz4nhcApbkfpAR4uARqCeSTht00IijKT28aVmfQvHNhR1fRNRVUmfTd6IWrxYzEQDCM0sbRPHJXSRPIzqLYI3E7djx/lRwZveX640iy35RpJJ+1yYj+mfLijgx7uTaSKPC9Iqum1E4CkCROVH2paz3KQuovAUdS5KmxjBI/9SNoNmcGPKxSSJs57vvLOFqFXpuzTAvSMTwLABaoOJ7nh5dMr9wrJOuH0EJtck3CNMkK++B76dURtF7aRuYga2ogZ/Mg+nBkXPuJP0PFKVfrj7YfemQNJVuWVV/LWbTesZbRup1cOR4WbdF/sFq0Qi2e0IUEXnV9hXK+JjVsuYHvvvec1LlLiwWRgxK9WJAkcS6JoD1Zj94cH7Kx5DiNCUncaEna8SpsX5EY0hI39pepQXXTdJZL4MFBCRxGbkrSdgI4zsI83cM/UGySGhkXxonRaDSNk1Pjds8w4N/5KnQR1fBIQmIfevkf5BeQxriBnQYADCMm6SoODZr6COl05nG0aDhpNAOgltFYUIKLjht55BWK8lHasJrNKWBv9qqL+Q7EkPa6RUGM1dlNShLg817lZUTmSBuphOSDwQuitBqUfGe2msPYzlF4BRnONq6fhf8tPNNi7sQxYzkgqfEElWnhLF9///gpaOZUrc0lSZ+AMvrhKlolFLhx7QQr0mJqCJiIltUQKc6chFzQ3tAyovk8ISnjw58bDNU4PjGsDMMQ4KE61pSnMswiZWOQICECkdMTw64gInLWCUh4mV4ZbaNbIm13ZOIZMSZiN01yiqzIh0ZDT9puTnV8vHDSqw7IHWSWE2uWuCjK+dqwOT9CS8dKhc6KIs7PrHNkygYWGLmmccDJG1VIxoFhc0SxdVhhl3WF2fctzNYXNqsrrHvfwrpqYVz/z+KWcdkyZuf6TnsTgunkPordxA+vXhOg3uDlvSc3dMz/O1f7wF++cEC3Y+fDCz9kf+M3J/G9s8zUMqeepA7MZ2/QPvF4EQ3ASJ1VocGo2V/5yXd+CBNDg2b94x9MhWCqaqybxiEiGMeGTdr9Ai+v4DrTLEWbcwBKKwHLDIk9pMQe5jD4hwIE0WWjXOoBx/49TqFR+Ncy+tBYM4CW0W02c93eCFqc931JjsaJNAZgPpOnkqGIHzWGyb+ysio8I6vjTIXMay1lZKMobeUMauGsG1m7c4aa04qaHucQOZf/pr9Yt4PWwj5jnOwpLU/Wy0auArIMoBsW6g26k4NlVcubPi8DRJWVdMg6siyWQ0VwbPIAwcJYC7i6LpN3PewdU0F1MixFdXJi0GNsS1aBfHjFrr4kOMqitsO8W9dZDznNjBbvWRkR7Wzy0a1VKH0myeJTp6oqIZSrR8Jo4YdgY8S5DvthQ1ABXbWVoY+LgLIgjHatLVSyjoLoUrNJTCltJzOsbUAu/J2aTSBXTLGlWSCN3qAh/SaNweJmc71su+WTSDZ1SVPM2/hy1nhwG2+MB7eX+L/ZpvlWO92gfR87CZhVl/crUYBAdyQEr+YGIDpduzuE/hnDEN0Z2YMu/L7E39aoi79nxSBVoJ0adndcyJ5XxqQOkKlV6QyEOjGmJFClouCfB6+cFBymEGwgaA/496YFbiJNyidKaOQGqqlPDT746zgD4UYWpB0cNBULLY4+ADwHPPPPMwXJyb1j5N4BOYDNSb0TSfHJJ/pw9u5cTIVC0nUH+X9N3LSB1sU74B3+8luG3RJsvkIjNyXVYox6/qWfUuP5VewvnBib6ozCmn+Z0z9mC37a89Fo1qc/5/PR3CL0Z9d1rK5Lf3rD7qg7pj8n/eFo5tGfY3cw7M/MFidIeqOR26U5M3fmddlPezQj7twEGComla83BFI8mbPRHP+h2A5xR2TMOZvNRpyHsTcfOzx1Mp4M6U93MLO8AfvZn7iTfs7ZfDQbeowdb+bNxoz9CfEcMsg528u5c0kQvAEvClgaTVmG4rSAZzL3L0s+iwcjzk8heQL42QhHdQ8ambotLYO7MH4CI5vvFb4MI9jKB/RMIbhxRqEFLeQKQrswcGlG7wP3qjEYfI1Bg6Y53dMoEhQFHXFEuWE/8N/mtJ7myFJoljoWp4tGcE6bfZzlenpmtYzi3/OWlGHTVLuc8UkwzpslP04WegcjFDANQJXd3OE01S7vcrNfsl5KpDIoabbxk5fOS+Y8NsWOXpK4OPDdtfk05HBS6A4GVB7wd1OgfK9GBBHbgojtc3H+rWyWvGHsUlPWY1WVdd5UTPDCNYc2fAbtAI7MjSRo1oLYVCAXTWAB/PGWUWq7skh10/Snaq7amfnjW46J2L5789wvS1tWfcu179V01p+i6QqJ6hqhqvvU98f6rmpv6XRqE/GJy0PHl89hfMibyvZOBneKUTDF7rlvP71rc9+1we/d5Pdu9NoOti3TFvqffSfM6sxz0eosem/enMeGdefmtP58zWnVyd26d3PemaymOcWmK8yWrIGbctvmPv7dbQ9tgSUSUh7+kVVAo00Zp4JKSU23ube+KF7WNkPSMJcOOAspjw6bpfDbR1mUcq/Kss/oWHsump0cBPhZARswtRLvi/X56a1PdSKsiJvoZJSHbVo0btMSAjd3b4KdlZP63rOoUi/RiX3lxGny+OYpguaOOZWL6HCh/Ebn5RRotm7L6OlzQLCDrRBDDtHHv3ugQFshBxySYvTvg9HDvwcfhwl/D1FfCu9TjDH74TX68SDQuQOtJHZWZgV/bdjGV0o8srCecuw0XpEtWhiSSyf1r0m+hHhcrHBmMAvnEsztlSctQag2mxCvwVXfHKeD4zEdipudZBn4acM0FVOPIWWLlcclxeI5OqMBQS+WCAtlqnhnEuFzcYxHKS9jukPiIiZL4qTJRTTH1epVEEjzuCbwJ5GdGgcHvjrn8R6epMBNy0h8j3AeOJeMZSEgWJIhAL6k23S4dAC2KQPnkqNDd7kyKk/lYUUTXGMiU9gydhNVxXxdVy4PNt6lVAqqFKu1D+hYWLIOihmEN04xi5RaS/xdaSEoxlDWn5r1Vkt5Fi2ZSlWTnzCunynVO9PbfoWTp83QGITnFdXVNLBYf+3AI5pFYrCLAX+xND6lpaHdK7KzGHYXb8VSCpoRbCnuUeg9Cz3fJUlDjWT7LB1/JKAnFFXZgJQN6Gd8SKCDcLaWBEgGjAQ67EyzIOfMRDvGPKeWDMw6JHYCU1Q2WkZnuUquMgTKqKmNPZVJqhZ7xnq+aYZX84yjOms/Mc/P5YkvAz4xOFTy3l9e0HHIVJZ6BHbfPrjVgG+O5GQSepD4Vu+N84KP71ouw9NR3YZjtHVl4c4kXK+q53uX9mCNpy6alRZEyuxRxDvUyGwsnOQ98YxolTbNe7F5EUTR+9WyxG22fmN8843xFcf1L8MoRgeRjpY1rVPNV7k6TFWT1SxJwUajfTdXQcbbBV2rNs8VdzHjVEat8hxLHK7C92H0IZTYq7AaBDyxsKp5aRfZYxfUiR6tO3237SCKvs+ebu86Oc379ICyXt2p2ViRCuu7ttrWNqvvIRXttalyRMTijv96aipORRSQDonjKG6YvzBexLHf5POIdmcX3wXACngX+WHme0gbTB9BI79ZEre0SnsBxp9z80uY+sHf2Ybvhkcw9ke3J7cMh4JlshPmL7Y93E9vfpolJL6me3f4NlgSJ4SiZVmNBpi4sU+SfB9z1lw8/cw67/gC4mssHfTPKk8+fCf2Cyd+j7vuTwyB3c7vKxLfvAFz1k2j+FEQNEx157YoeVa3hrgqkXlCJMD+ohQmaxCAdGKyiK5Jo6nT5JKAOp6fQB1CtDzUpmwZt5t8pzDUIkkfheA2IH/fxc6CCFvA9bQj9kNsvMyO0W/rnq38wHvEN5l/51+uYqXlXRoqyeq8TU9k/i52pS6zKOrlHfnLLKd5xCLYGOXBL3GrK91S/9hJyLBfAAmJJdinLE4kgdI0EZJOWi9gXlIpKxkiDh4/wIiBR/egc3ghUYRdZ31XAC3SRMgbDeSNFjIJYALwNOBKhogjh9AKFFfZBiKYuP4ahsFXJMY9IAWClKxIckW+owY2Lj48E6xfSapVQNO9YpDCTkzbq9hrJTYpTJVzdsZAmCYpABtR1BMPgpaUt1AohP1wK1l2JKGSaIaZXqFjgqPqt2xyWIXJarmM4hRsII9O+83yXnWqdxdoKcmFskMiilY2RaHlkosC9HD5PnhMcFdxDGO1nJiQJfVgLNGFUeJI+UbdQmU7WeThRk1qKhvMsAA2afLyi8B/zg+yepB959v6Kf9S4katp1zjNEqd4EkUJEq9o+BXPCVF62mfFxmsOiDVkg+nCKBUb+0mugQAAKcAFuJmPIiJ23EhpwOmExhU9Ce1rzKwgiMmssIuEetW/H4IJA+oiKEs3OYfSmco8pozajmepHC55F5HHxTJgeb+QPzLq7QkuptdRXdzF9HdfITobupFxytX/N4iuqLqguwQsVmhinzKcTGo9gZH4J+oaYuF3m6q5aMM1jsIScEQJFUu/IzKKQTj41w4ALSncu064LYl4rpdw4vc1QI6XseNiZOSbwOCXw2TgZr5GSr62aGnEHHjd6GaB0YXz0ZkWw8l8Csq2RyetoceXhDqE4qLi6RknUq8cqp8CRVyG2bXoyxqBgm+3ZoeTPmR78qWgsRFwB6MT5zesAll00hYAslmztfF7ncx3ly3bqXdnKdF4TzLtivf5A8FluAXwlkAYWd/aWs439qvbgsvLS6zDTH0+MejWaIrMc+UdyXkiM66BpFmlhjWN9QJk295IaPUDHw3O91vjALB/fa4c1sVjn7xnhuQdPf7i4xkzdmm4o+O7xaW2ipE2CqE0jKsjj1oTnetjsQTJh7g6b3D4lDe1kg4XyzDE0Ia1fFDTRPeRc/YyaOFeEpnW6WszmCHFqlp4TZWh7Ywcsq+tkRg9jRHO8B4f6YenxRiunQABo9SGYGlwxY5iYMT0UMQB+KH9QP1VOWQuxDFfCKYI6iBbGyylA0VNwX8TT08H8aKqd/AM9UBgdE5xhNLt8LiCbcKGGyeWgku/RQcBMCnhvNZLq/zstJlUcYTyRrmA4DxV7p4YxwZX31VZFfQ02xjF+MHytxyh+3te6VFO0lHJR3U/2TGlHfthC55Eq3CVNS9+xpV3CLKdAvsmwOhcXMDB1MLO4eCZQaRDFup1ZLBJqtvMSKIfJyeCIYarhWXFmGVT1k25WrI1VRkJ/Pe3so7dE2JO5UWygyoNCqE1qzHr1h1Vj5nYG+9r1433ijDraI7J7k4tJsRlNJKZYnzhrFX3rIhjQPrarW9jxNF/Z8alcUh6uCk8Gkq9bVSW9f12oriXMu6uq7TVeljvUVP13VaWqmja62Orqt1DIWEGqqXUrMOWaueO6uLZL6uVaVc1ynlXmUJVZO1/JdmCO8Efkh+5U6JPa0BTNI4ek8qVvKrKD9xlgic/L5yYrIV+j8iamqZC1zdNT/rFLtXO7Hxyuo2nLBV9PycWBtNspZm9NsGk+c+FHXQPjf++lc0U3FjQg2GMK5WoVRNqbp51K6cR+0v8+iXebRmHj39dPPo3m6Tp10xedpfJs8/9eR5+gkmT/p/MRhWXKdCUilO0QjJh+y3vFlJyMAJXRfkaPJ17IyFqkhbQemOUTLcb7I9MiZG7kprtGpEMJcE25fwCpe18LKDm6ozDtRPFq4josEHGoDJJjK84sOeao7Xi1g3FViyQHCTM95dg/j8J6aeih25iPEXlyidA+SePNyeSlOUEN0WsDRX4YixzHorLDs1l6yCVIh5f7LISrY4zIgU8X/ooDuFUwzOmxwSp/SmdwzfSEgbOYqC2o75/LKp0kUKvhy/KaSP4/wpc5ba7XLF6xaXchjKrLF9mOdBvchbBatEhqdXP+U49EvEU+tXfH6d0ZuWzrpGsafc4SUiHnK+S3v0i0vIkMCBZs0n6yiYz/sHQzotFovk4VKj1uL4+/lVBHqeKEL41KrIzf1V5GYnFdlqr6o6IiLUKkm5hvdTEgnxz6Mk2SVrany0ZWjDnC3OThG0PK9doCzuw8R4vnBJULEBC0OXofcEHE6vchnQ86/NpnyRoh9SqsK6nLqiwgq+O31KWUJjK37ZAl9Bmt7ema9GmsKtqOZUC5qvRO4Ai/cCf0evBabefHYxcAV0GjthgjvNf4r9SxYASKMljAHzKvowcb2KoyWJ05uG6S+cS9KOCWq/H17i/S50zw0IyTObOxBot7PbR9t/RNECCdg7IuJVY216C3IClgnF7C/XZiZu1hyyqN+6TuA2rp24oZSLVvODW3GZeLNcZ0cCRUp5S+xGioFX0MrupqWnCFBmYB6YmbIw/J0bSQSX5EQvSUbh6Fg1y4vPqtD5Bb/Poay81g9uSwP/ht7e2BmQBY6zUN+swhX0fo6WArn1TuRKuxiC4LkzI4G4uYOvKtH0V8q5Bz634Pr7AvoFhaEhNbp0HuCntH5OUzINwpuv3mBHQukvhXiaCMWuSX5Mr9pGOKtjQw2EC5O1WKjD7N4jYSgyqWjwzjabLss2rM4QB9dyIzZR41iG3CWa2VmUzeffM0Troxyprh81dZIQe0R2X7NWZIGojtWdkHstQv9T24u3FCNkW9bXoG8Pbn3UP6p+ejw+ngi13caJaroWEUoNcxgFfcIG+cyJoBk6VhD2EV6/jjITRgRDq+hcy3OPtjRlyfl5D2O7i8TMzWffS/XfrFLSqBlHKZDHYbM9sTxyaWppa8ZlrlExDvvaYmJp/thZ43Sa2jbur/l31uCyYV2rwhrllGWyg3buyRY0qOrze6tEgb61dwmgO6hQCXpWjP7VQJkW1MHA2J+tDZnY3tWQkurmStmOqFaiEktaqaiy2PsLyvcdIjQUhMmYwdxSoCOhHi08ZoVJlPamWbUDkAZeHkdrklRb8B/lIxQFdNzASZLnfpJ2wGIBSzecRyhK/gZDbjndJTzEwdgVh7XbRLkXxHYVv+FMKYqvrCQIrO9SbXyTxBSucFbK4jQab/8rfHBb9JHN2VtlAw++flFirapQ+myG0CsZMo+iGiZ7ScNUcrOtOGUpcQC2R6w621nL4rYVgCQlS3FtQy8KJk6Gsk1q5rlhKnJienNPOTHkXE70URRTyawWEweoElOWXSkmDqCKKU8ue1hor/Yt9wqvVkn4BFHMEHXyZUTLsnMxcAo5T2C8wVdK6FsV5DqVDPaSLACgA1W9JClLKkIfim5VA+7JK9XV25b5/RIqJal7C3Ghivi/2uQw9nyLtzbgQERguMx0o6UTitqv7oNcHnQxihMAhU845HKKxfiTpXToRjoztz8W7DicWTq4jvd8w1CLp/jxRpvOgB6ot/EGGnGFQK7BT0s+jUv852XvwnpESYgDp0w816j80nDxvAyaU+WBYMdhoEDWDgRCNuvppq3JoZ3c7GqyePc2nfCmnJnVSnOgSiRB0kdpGvszmMwbJm3hlti0SiBuHuEzW58uwMcpamduBaTeqc+gstH8f0T4clZbJcKzCwHo880D6s/TY8nzqFmOY+RvNFV2sLIRtpNAMsJMIi9hoMJ6ZqnmVOVBmcg/HQ84GJV5wMuwFKC7hevK7odCaOdoXZkSt8DyF5w+YRw6o1llYuYAcnhOuGN9/YI+6gYThzLSYQAIN07RWD+2F8xDj/E9Lz+8fBL4wMxr6UywsOQTglv0OhPXg9uMEndV2jlpGnsBQb3VrghlMQDB5q/YCZKbuFmYTsCRl3g4SIf6CIoDlXOtgRdjrXoXCaUtHtPGWlA7l8tQIbVFmGwJJ8fPtO+0EB1LUW7e2sauYufL3LNARYMGKZoGzcLge6NtdQZof9mdLlkoN0R8TA0NWU94/LakLIUQuLbgwSn3SrxkT7MX5d6kLTlwou4qkfR8vgqyyme06zS58G+36TEnmiPgdp6t8hS4kU9NzdOj4tmODIpWuZUVRr+k1zjSaKnDguQCCT4kHCrXo+L0bI5FMwo8+ilhZk9KllFZToHLvgXkjeZxlrplB0EPaOi3bUgiEYejugWHUriLvYZJI3B5Abi0ItCHTym+VUFfGDplCYpDZ8xXg7YRy1cQRGo8iCOQYykFvY1yKh0nJqCG5gvdMpQU+1qY1q6WMM0QzP8ujhZvuGWqP0Yo7Q7Aq068V2t2eQZfD6epjQ9+6EUfOh65Bg+DlkqBpJv5K2DwqStpc49ajG0cikXh57bi5Md/ljnAI+/dKkkXLAaoFFRJtfq5HTdJXq1/pm4ffaEkTth5/oZqQ0tyELGUussEbbU6iC2hP1Qq0TS+Ft7TOCldb/JxS7r6+y/vRNNZpZFZ0Ux0zRL1UKyhuAntobYlqzrVlgVn6GO0QOHWKtHCkpbQGH8fnNS9EvRW7TF8UwrxfPmWYa5WC0R/gbmNt42Y0IAyfaL0wW2F4m28JQ5BSlgk4G59qXyDlc3vjFFCAO4VdRpbObowG+nGA3G2UiouxC2yH6zgcuChqtRNyYmsDlrIlawZ3PLb4wKPrTML5NlVTpCcgdbXeU9qucSNoyB4LNlkt3TS1ZWAbxlSDlrGjFw51z6+4mri/UpOmJobtQzRkablPAvT6O8++dC41aADzSBy30NKSBxQooLgRj70XpbnIlolVDNQpmoMLbv6CPeZguTEDacY2mKV+8+WUXz8Jo1pGabmBHTujrJuyV/qpYHfKCTC9ZGKF1cFWHUTdH4a8uJauSMx33yU12yXQqmgq9lja9GSlb7F3aSbZumB6Io1tFIZqWK9bC+CHvjYXgQ2h7wgpr92Uiq7WsA1V+Y0pzsogipqITgBphFADNU8tJiIly2U5ZSZarLkg5yAHvmHLEqgYPN0HTq9UKzoDtDTRT5ODcv45htJZCLwgQLMXC+BY9lBVKTFtz6Y0mW3Cgz3jhXPcbelY7laWdUV3dE4b1o2tRzsVPRBVdFVlZUksqloqd/ylvohc85rmuq3vKlyaLGtftB48wpvtN/Wt1W2XvyRjfXbxzRWeXi5Q1v9doe2KhbHzcojIDtNXxGzCNTZa8cpZqcJZisjNE63bQ59nR01uOdMmuGXp9PyqRgEuP/keFHMiqzQqXqd6sV1+frXei70N3vkAVB55jGfYAbxjjBYrmjjlsuQ1cqUrjvR/dltsjLk57L0R4rkE1XbZFJxEKuEpj4KoL8hZCdd/XBFSKDT1Wx8dIIUyi0HtjwSpM5v6GIw3GaHpRRcITa/W/4pmTugOGpItHiIumS0N/HhhSKf0VYevam19EWZlAFxX7vCLncNYNZAuOfAd7Cjy57fTyciSoSm6kHoMkenRn9b7donRl/VP6nU4wp28Tab0nU2MrvS54F8glF31K+yHsdGe2tFDrZV5LSqIn54p4q0t1dEs5gpktjdBy7/uN/SS9mpvFUjubRjuhT3NzHuhyFc41DR4YfKmNdQJCRBS3FcFmwWCvtPdbnnU5am82WL9tLcVrt0ksS/Jkfs9ZaNtCam20VR14DaAEbpwloWqPopJE/5mz2f7rpa+XYh8aJRzeZXFrzgQFmKBi4PF4ugLHEqX89UvkJv90v0StfoyZfgla/NK+UzFlJ1t1j1xXjle5HoaxLs0QY++eNdE6AluCXiSHhlCTSHz7rKydj8YNE9dpFWHuGxXWdpVkEVp3dqwRY04sja3LQ6owFZVMIKFqsf4u0dbdlzLq0r6I+a6IHROJ4H9DCwya5/roatP22kkVjlOQoVEmuVe+ZmtzMwd4rvbgtfV9V5lWKBtJGWa+XIxl7d2bD8kJfUCWr1QwSpUg4RpmID8p6y6PwzTgulS8wx59ug5IvvqPFVtc7IisYRT1Lr3hX3LpfAtIdfDAzKYxcAVzNfzrTpPi+YibLKZjcNbMjibSX92s3bJegsuNIe1PFceLG2XQcnbX93SSherFMuWtJ3Wwcou0vqhLBR9YHNAH8elWD1/bdpBY/1/NtVIpv4Za2o2rKUt2X5Dn9+xveWvxWg2ERH5SRmx2mu0z/SJSL0Zrq3aVIrLL3yqQnwOorSl5FHGs3OVZSk4G3Ow6STBaCyGxnh5/T4ECxnf5meHh+mMSGJC1NAO16F7SsSk9Nj3OBu0B1WJ/vzKPDwKY+LECjvnx5TOZ0e04UlAy2Hk333irjvoQ77WpyLNLq8DBD1kCLJ5OltHxfObAZe8H5d1jdBOn3nrDtUEAa/Fr7Rb6FP+K9//s+21bHGxv/9P1ZnaJz965//G74nkPmvf/4vu2OfG2EU/kHi6MgeQtWhjNP8L14fsWBw+MMLsoYUj3gKU3iiG1TGu3BRBaAZa/MhE/dp8rxM/Be54ijIyisY+6evs7mYNV2n08n41jYbbVPeKKCoge9SpT2M3JSk7QRwnMX+af6yFVN8qj3sC5Vnmgcv5vRC/1rlmjKwDvDwJorC7P0LTiJd8vAAjovSox8NMyWLJRoWSIfEMQgMTDg8GlIYm+z9gv9489PLDnV1G0iww09wq/RY3c2m1KuBGqLk738Yco/o6F6vYOwWT19gr97SXbK6ZM0oVWm/pl3eJRHoz+0+7l7cPzL2n7IAWX4FDtj3GNYDj9Zw6FJ9s2P8gCbmIUvHTVfs3QIaEtxvGfvCawVI8fsfJy+vR/ZNHPftbvjk8LeX375/9WNMrn4e/nLwbHw4+LH//FX35s31aBAv+i8Ol48+rK77y+X3z0fv/tZ9P1mk1/51MHa+Hx2Ef3vZP3i0fP7q+vLkJC+KPqiAJfEeiRnK4wWY++hv3z4u/8foCLcnAOjtvmPDXzYWYMGPPngj+8XpWUg5u93HgwNIFkAhl3Zh+glQhg04mEpPDUAquDv70I0oqXP4fVNFyyrTsrbQUo4rIUFMlq8wopXKxFRcdoR0YSKGxDYMUd3+aNQbWN1efzIe9AaY56xRDrosHkiFbObB7efXTiOPHcu2rMFkOO4P7P5kMulNOAS9jbqKZvYsM9bhbDJkBwJ69L1Yu9dnn32bfg5HkDPE/+jnGH6O8OVb9jkZA7CFT/SxZ3Ut/EYebJu9omtjwgjfgRywhC4kdK0eJEzYe7V9PJqAJXR7DMLq4uEE+j9WyBDRJ/hsoGXRBBuPMIwsmooJowGyjCWNKUYPf3Z7+TcSwwmCg2OJQ6z0kL6I2wNqfaRo2+fnqIDS5n0Q0YipeVUMnGlW1ub80b/9DIeqF3aefYzR1IKdySD0PhHa90DlCqUtnvFCzcwSMkXdVgb83lKI/fGFnG+v63+FBoz3MhzL3JxvxCGYD7fwEybXU3VCZ3/taGMoU/0/jKd0+9CR8eTVL4bFiX1zmU71pTAGcIqnVicM/Sy4cIQGB9k3fA+NofginxF6xPNGE8/pOpbV99yu4zrDkTebdOeD8Wjg9Dl7n8kEMygY8e5pirnRAp2aizr7iOYBeOAsE8xEs03XJgqNvxgc6JOYeLUMVTfmZzYFabUaZ2fA99jq9QfdCYwi8NHrWyN7Yvfox6RvjbuTodEyhIG8eGL1DFK7Q7s/7PdhACtGcgY/6A+GvVGPkx317WHfMERku2N1x/bYQngAGY4m47HVZx8wWo77ls1LHgwn/clAeNwV2e73e9bYHo4ZsgUqa+O0APB9qMR4ZDBKfWvQt60xvpXOHSiqpLxBmxXS/7iONXdJl/R67mhgOX1vMJ4MJt7MHdij2ZDMB5Pu5+1Yn9+36dlU8wfjTPXHgy/ezRfv5s/t3bz8OTm4Hr/74dfnh5P1j48Oxj/++PN1fJMexv3J73/8eDhKHz27dp+ur/44fPrC/v1wsLb//ioEjygehqOf+4evb66fvEp+/eHlu9Hhu6vkMI2+v371xbv55N6NNexawwFMVQOw0+3xWPBuNFnbvJsezDXd3ghmn+HE6g3QOVHcGw3RL+7NF/fmi3vzud2b4dgd9SZO13InVt+Z92bd8bzbc8f2zLO67sz94t58Givv/3sHZwSewsSaMH9g3EdHo888i57dt6wuz+mOx71Bd2AobsYEXJdJt8/cjMlo0h1wH4WP/WPq+vBRXkEeD2BU7fZYYTA62106AVCHCFyXgcFyRv3B2O6NDdm16kFZvWG/R0Hs0bibUxpao+HAGvE69Luj8RDY/lgH56PtXHZdyjOP3tVYmGajeXfmEduaDCzSHzi9iQsyH8LcMPMG9nzgUMNk55iECLzVzxKBtw4H51NujkduDBa41iDPbGaUFQfi4IqVvN/JxXmB2fsl1CdBRPcpZN8dFxNoMfTZlWmRExMY6Fzyq59eNST0nOg8di75IUhl8fcp//yOQ2AdMmhpATEjjPma89N54yrnp9mDrtkBx5wyfPJl58c3z4B4ho0vG0zZU8DRKnbJU3pzSIUM/4Lj+b5xYCjo/F4eUSw5tc7cj5OsbFozQChyCwdlI7eEVsjb/JH/B2mC0aw=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_16a3675ef842438cb85db5d1411974db\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
     "## Save and restore as pure dictionaries\n",
     "\n",
-    "You might prefer to work with Python built-in container types when interacting with checkpoint libraries. In that case, you can use the `nnx.State.to_pure_dict` and `nnx.State.replace_by_pure_dict` API to convert an `nnx.State` to and from pure nested dictionaries."
+    "When interacting with checkpoint libraries (like Orbax), you may prefer to work with Python built-in container types. In this case, you can use the [`nnx.State.to_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L170) and [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179) API to convert an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) to and from pure nested dictionaries."
    ]
   },
   {
@@ -253,24 +293,34 @@
     "nnx.display(pure_dict_state)\n",
     "checkpointer.save(ckpt_dir / 'pure_dict', pure_dict_state)\n",
     "\n",
-    "# Restore as pure dict\n",
+    "# Restore as a pure dictionary.\n",
     "restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\n",
     "abstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\n",
     "graphdef, abstract_state = nnx.split(abstract_model)\n",
     "abstract_state.replace_by_pure_dict(restored_pure_dict)\n",
     "model = nnx.merge(graphdef, abstract_state)\n",
-    "assert model(x).shape == (3, 4)  # the model still works!"
+    "assert model(x).shape == (3, 4)  # The model still works!"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_277cbde407034bf9addfa8e4969bce7e\" ></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_277cbde407034bf9addfa8e4969bce7e\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWQtT4sgW/is9maoBrhITnoJK3YC8nFFHcUbH3S23STpJS+jETgPilv/9nk6CPETH2XHn7guqBDqnz7PPOd9pd0Mx9UhNFZyQ0PQDcsV9X6DfUOCHVFCfVREnHhZ0THaQ7TORtfGQetMqGvrMDwNswvrEpYJkox9VFHBY8WgoshHrrJgGsMp8Bst9bA4c7o+YlTV9z+fVeOsOSn71PSAAftQSbhXZVAAZE4SJHRRgy6LMyXrEFlWUM10phJGsS6jjwoquFiUbJjAFnR+2JV+yYxrSPvWoAM3xSPgPtFnKBKcspGY2pHckfpqoe7+7Fbtn98E9WT5iIJPDWmhyGggk7dtL4SDwqImlx7Z8UxBpPSd4mKql05m9GjgU5IUCWcRmIdpDwqWh6hBxCt4+8i2SzqiuHwo1eg6mEYGuAsKkyYYpucpNP/2y7kkHM8sj8JiNPG8nlqCCmj3fZ7Canvh8kEGLOvjnsCQfLS0LasrFgHDb50PMTKIyf5LORPEFAelHT1A23rSL8rkM8KE2Sq9orXqEOcJFe3tIkyTPqs6JGHEGfkfEC8lcMXfEpGarrEOX2kLqFxHIL/fwfkJCGk4Vs/yJysnNiITCYHQYhavF8ZCkY59kJI+dR4KCUejGbtxZY+NMxF5sxjNWvlwHqUUcSOE7jhdn5VWUOXBaA8lLrhBPbCIyhgOeRFJqF/1WB2Qqna5wRSqUEKumh8PwAyRnwjetPPC8GsIxVGbC7zPgTzj+0Rmv7W6tSwCLjlHEcE9ZLh8KErgPlpLbPUVTIHW5eEziM1ARnMHg0XPJsN4DablnZrsCyRiXsaiOXJn+cAjr0QGK6spbHL2kbSskuMp8ka66/pjwzBr6hBxUhfhbCxw1eNn2nCKUVcRZFJnXi3pREoAZhHNiXQVQ3ojrexbhi4Rl+U6qa2RHFVGBoZzIzTaQ4z7Yz8AxT9RlsGqJ7Cp2GVBbNASh01n9XSVENeThPvGq1T6B1CYLWpnRa2etvLgGZ3VZhJNare3MZVEWVea+58ti/qTMyOePJVuYD0KCHTgu7PHuKFwrSy4O07WIZ22tH+IQmy4xB8TKZNB/MnMd5Nb1m2b0SxpGzaSKUj/nin0z9f9Ub3nTk0qWfoCSMo5S8IiHMoCBDy2V8DVyafh6YqNUiARloyIQPnXGX0fq3DxBbsVjKSoNr2zKQ3Hlsyt5/Nek1nOppOaKMpvWhgp9t/pxxFdVlFYNMXcABMVqRAl9/53SoB4G0/5ICEAf6wrQ/PG6Q6sgZYUKHAnIcj3xz0QvWMoKME0dYjgVFHuoNx32fS9ExyMh7bVQI94Jn8EUEiM7If0BgMy48g6hwbhQwwECMgHbKQ6J9QBN3xJNvnceH/N4d4QdNbVChqtWxvmxxor15W6+U53g8MqEPgCOfdiPbbHUPWZ1+jmZK3uWRS66Ho0xT2ezFhY4ixkENkInmcVlKURCLo7Z7DRHbJEeIgIeAzyd9Ufi20x50AACQ4n1ZlmTSCR6Q4eBzwVmj3j3uT8g7EquzIvR1727sG3Bn7Mw36sSwIBi1pUJ8NfihCWqLg8YwHOZ8JWq60LqJI10KVVN7JlpmIIAguvBbYTa1FBguf9B3z9Mkz5AIai7sSaWL8B2qcWi80IXS8QzwZxBSl3NSvbMy7aNTT2/hjAAcPvbw7DHk9lO1qXE/GQpq6lRwZxPlNVoHsQ863BsUQhIWs8XLeJsIh/OqkOQBulZMt3N+OwCKJWVIFpCif8eqfKoZL6aS5eCOTPmXp240Yn0PBxA8fk6gPv26vy0hBi2RTIiInILWWo9SfMaeqwTkRgaE8wnnnWuWB4p1OVRBj3DYdVUtjRcrSd8ibAn+bzifYIc8dAbg3M8VW3uD2FsNEdyllFlIwrVMfZgtkxnMmrow1AZtSc5HMpPNe6RcjB8YZdUUiiDMg+jeOgSIuS8Tiao0ev1pDU9uSan7+ghzLbRmNObMjP963+TzmzKjvc7u/TiXMTk/YOXrE2Sq5+CBiUn5GYVjbiXli2jKp9vTXzbzu30oSmVCpuWVmkfOkbdiF7dE8Pwo2/10wn87bQMo2k896oPDcMZ+O+tbrPemHwxjLMvjQPjsFtvGC3nttv54IqwfkiJk2/tX+Q+dEtfxr1gRD8eFs/0g4vu6efD8fnhnfg4bbUaG+fO4IzW9zWX7p+MDppW+1rr9LfscdcKbt6X3JtzSk9Gh6ztduxPwvhUqh/xgtHqskGzZH4ajdjGafHGDAeTsd3ytm5unaa/7fQPJu1tvWNsMeO0+IHzA/10w7nTTi3NOLB156jcmLSvc47mT0en5fKwqZcmnYvKseME5GwwLZBu/65o9vlxW2DDOekeTfZxOA1PRt3uxXmzNTE+ngTdL9anra0Np3xWvsgLzX7/8cYYF4HnB+OobBxOjKFzd9rbGF32SPPiNmeXzLujwmlnWhzVjfd39eugFeRp56TR1C5HHwu9MrPrH5qd1uHQoBvb42bOZbpb3uh/nlxcTzp8vN/+1GDXdrPpiI1j89LzysVK42BS33YrhcPDdi/fvjScYbd4XT+piLM26VSa9Xq3nd93CqdbX8xp32hDTD+/3zJO2tgghw3P6Nw1j51L4ZTqH53j4+5+fUBPiqRVv2jUWybVApf7AYOzEVw29/U7fdCzG7Zwp+9Zx8KtsGNrR8N286hUt4ybz58DLMLe5dCyMK3k7LtK4RO9vikFQ1469r80epS3h+ODdr533su3mjmzfmKfbXQ8P2gXWuGkiJ2b0ja9JL0jLzhn9U6XWIecjM5v2o2hft7ig17vtpgrnZ+HEwM0yqDo7k2kU9GxTsmW9Sv8ech+bPkBNOJ5SkY3hqqqPkOxGefsL8Dr+TsYN7rCirBSDOOANxwPZqJ0jKaWLxghBc98mb5AlqAtuRZCeZAsJHyUmAtPMBWI4TF1sPC5CpyDvo+5pU44FeQMJq30nBcYm/Ca32IBTkgrC9hS3l+BlDM6JABC07MLzkf7OBkCUHy09X4T5TRNi+AAFF9ABuloSlovdwFAKnPl5Hw4q2Dyyk9Bb1ELUw8Km/CRJH4TVTaAAGyEPajGFHxGsCXx78ai75K7uK/cwkksPbuGW75oWQUvSm037tC7lAWjpNMoUU/u+7fKWiZJ+4aHcesGJaLNtd92txJmiwo8wszKc89XHi5epCm1d29vc+WdGNPp8Q/wBJDXquhpo7/HxudtS81HhxTyWUOemr3UN6ZNdG2ZSaGHuWZP+SmV2Jj6RUFRm9tTFgafKnp3M/LFzgJZvADj6+qgBl01OjhQmFz4nri3lrht9vGHxqsPU/NzwVq6TlVqR4CLEsLNJbplmKrU0DrtlzCk8qOj9VNKGvviqM3IvzN6L43lC2I1IJwR77loLSPoGUtq7SkPd9rYrJRxv1wh+VylULTzFd3WyxVNK1cqhGwXV/VYdxeu1K7xrRohWmR7Phb5XLqwiQqZr9v8GHWvVbKSL1swDfY1oucKpl3CBV3D2xbRSnYu39fzSu0Prihri+V6X7zzxM4T/njniJ0XH4eV5Nj8xgR6IvY/MLvi4/kN+TXb8Ir1cemfRVCFoJd/piG0bXoXcUQutSDroIWj+dwnB8LfdWb/Ac595uP+r9oIfkzbfsXCH6uU+weAqtzLopP7F1T9GXIp922gKve3BVXE1i29tJ2HQV8rFPM6IKuyTQr2dsXMl4p9/U8BqvSSXsnlynqltL1dsLXt7XLeKtsAsHLlMraK2r+g6s8FqnLf2vdz/4Kqv4ZzXwCqXvMU/22a9u912MMOi45r/wOOvCUk</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_277cbde407034bf9addfa8e4969bce7e\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
+    "\n",
+    "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet620aS6H89BaI4IWmRFMA7RUnz+ZKLN7bjsZPJZrX6ZJBoSrBBgAFAmYqG/8+8xzkPcF7hPMo8yanqC9DdaICU7OzsfrEnYxPdVdXV1dXdVdW3Y8+/tpL0JiAn+56fLAP35sgKo5DsW753sj+P4guPzEkcE+9i3B16Trc/tYnT6c3mA7fn2O7II/Zg3ulOne7+6XGydEP4G+mdtt04dm+u/d8vZlGYun5IYuvW+nDlp6QFcDOCBcULN5hYG8sE3PbDeQQoc0hpzd2FHwBviyiMKPbEmkVBFB9ZX7r0z8RauPGlH7amUZpGiyPLbnf6ZDFRS1zGpLo4P1yu0rP0ZgkCid3wkuyfAwvXJE79mRu03MC/DIEL3/MCoDT3g5QAD5dALYF8UncaVgRF+elN3W73G3cu7OgquqaCKpK+G71wtZiSGAiGUVo/mkezVdIAstMo9kjcil3PXyVHVne5/jiS7DdlGsmLNhnSPxNe3JHlLNdWEgW+l2dVlNpOAJLEia4vVa1HWUj9JeAoijyxllHip34EzeZOgYdVCmlTd/b+Mo5WodfiLNOCTAxPA4AFKq7n+eEl06vZFZL1Q2ihFrkmYZqIwj74Xnp1BK2XtpA5yJpYyNk8iD4cWdd+4k9RcYrV+r3lhx5ZQ8m2bVfXchqtd6xltG4lV66HRdv0f1gtWqEmT+hAAq+6uUIZX+MKtmaBP3vvual7lxYLIhclerEgSeJeEkl7RI/etGnKhTudxuRaAhiMOi5l5/iQDTbHaUxIMouWpBWvwtYViSEtmcX+MrWo8tbc5RKYdFFEh9EsJWkrARx3UTvdwz/AV5Jagk3rxKrXG9bJqXW7Z1nw33wVzhDV8khCYh+Ggd/JzyCuUR17FQBYVkzSVRxaNPUR0mnP42hRd9NoCkBNq76gBBdQJ4+8Qlk/Sut2ozEB7M1eeTHfgpzSbicviLE6vUlJAnzeqzxBZI60kUpIPli8IEqrTsm3p6s5DP4chVeQ4Wzj+ln4X8IzLeZOHDOWA5JaT1CZFu7y9XePn4LqTvTaXJL0CWirH66iVUKB69dusCJNpoaAiWiihkhx6ibkgnaXphXN5wlJGR/+3GKo1vGJZQsMS4KH6tgTnsow85SNRYKESEROTyynhIjMWTsg4WV6ZbWsToG001aJC2JMxLM0ySiyIh9adTNppzEx8fHCTa/aIHeQWUasUeAiL+cry+H8SC0daxU6y4s4P7PPkSkHWGDkGtYBJ2+VIVkHlsMR5dZhhV1WFebctzDHXNi0qrDOfQvr6IVx/T+Lm9Zl05qemzvtTQi21exRPEv88Oo1Aep1Xt57ckMnhb9xtQ/85QsXdDt2P7zwQ/YvfnMS37lLoZYZ9SR1YcJ7gwaMx4uoA0bqrnINRs3+wk++9UOYOeo06+9/ZyoEc1l93bAOEcE6thzS6uV4WQXXQrM0bc4AKK0ETDck9pASe5jB4B8KEESX9WKpBxz7tziFRuFfy+hDfc0Amlan0ch0eyNpcdb3FTlaJ8oYgPlMnlqGJn7UGCb/0srq8IysiTMdMqu1kiFGUdrKAmrhruui3TlDjUlJTY8ziIzLf9E/rNtBa2GfsU72tJYn62U9UwFVBtANc/UG3cnARNWyps/KAFGJkg5ZR1bFcqgJjk0eIFgYawHX1GWyroe9YyKpjsDSVCcjBj3GsVUVyIZX7OpLgqMsajvMu1Wd9ZDTFLR4zxJEjLPJR7dWrvRCkvmnSVV1QihXj4TRwg/BxogzHfbDuqQCpmprQx8XAWVBGu2aW6iIjoLoSrMpTGltpzJsbEAu/J2aTSKXT7GFWSCN3qAh/SaNwSRnc71qu2WTiJi6lCnmbXw5rT+4jTfWg9tL/Gu6abw1TjfoAMRuAmbV5f1KlCDQXwnB7bkBiHbH6Qygf8YwRLeHTr8Dvy/xtz3s4O9pPkjlaKeW0xnlsueVqVEPqWZUaQFCvZyaIlCtouDAB6/cFDyqEGwgaA/476YJfiRNyiZKaOQ6qqlPDT7451iAcCML0g4OGpqFFkcfAJ4DnvnnQkEycu8YuXdADmAzUu9kUnzyiT6cvTuXU6GQdN1G/l+TWVpH6+Id8A7/+E3LaUo2X66Rm4JqMUY9/9JPqfH8KvYXboxNdUZha1/O6Z9aE3468+Fw2qM/5/Ph3Cb0Z2fm2p0Z/ekNOsPOiP4c9wbDqUd/jmb9QW9aa3KCpDsczjo0Zzqbeh320xlOyWxeAxgqJp2vNwRSPJWz4Rz/R7FdMhuSEedsOh1yHkbefOTy1PFoPKA/Z/2p7fXZz954Nu5lnM2H04HH2PGm3nTE2B8TzyX9jLO9jLsZCYI34EUBS8MJy9CcFvBM5v5lwWfxYMT5MSRPAF+McFT3oJGp29K0uAvjJzCy+V7uyzCCzWxAFwrBjTMKLWkhVxDahYHLWvQ+mF3V+/2vMKrQqE32DIoERUFHHFJu2A/8rzGppjm0NZqFjsXpohGc0WYfZ5mentlNK//vvKlkODTVKWZ8EozzRsGPU4XexggFTANQ5VnmcNb0Lj/jZr9ivRRICShltvGTl+5L5jw25I5ekLg88N21+QzkcFLo9PtUHvBvQ6J8r0YEETuSiJ1zef4tbZasYZxCU1ZjlZV13tBM8Nw1hzZ8Bu0AjsyNImjWgthUIBdDYAH88aZVaLuiSE3T9KdqrsqZ+eNbjonYuXvz3C/LWFZ1y7Xu1XT2n6LpcomaGqGs+1T3x+qu6mzpdHoT8YnLQ8eXz2F8yJuo9o6AO8UomGb33Lef3rW579rg927yezd6ZQfblulI/c+5E2Z55rlsdea9N2vOY8u+c3Paf77mtKvkbt+7Oe9M1tCcctPlZoto4IbatpmPf3fbw1hggYSSh39UFTBok+BUUiml6Tb31hfNy9pmSFq1pQvOQsqjw7VC+O2jLEq1V4nsMzrWnstmJwcBflbABkytxPtsfX5661OfCEviJiYZZWGbJo3bNKXAzd2bYGflpL73NCrVS3RiX7lxmjy+eYqgmWNO5SI7XCi/4XkxBZqt07S65hwQbH8rxIBD9PDfLijQVsg+h6QYvftgdPHf/sdhwr8D1Jfc+5RjzH54jX48CHTuQivJnZVZwV9ZjvWFFo/MracMO41XZIsWhuTSTf1rki0hHucrnAJm4V6Cub3ylCUI3WaT4jW46pvhtHE8pkNxo50sAz+t12qaqceQxGLlcUGxeI7JaEDQiyXCQpk63plC+Fwe41HKy5huobiIyZK4aXIRzXG1ehUEyjxuCPwpZCfWwYGvz3m8hycpcNO0Et8jnAfOJWNZCggWZAiAL+k+Hi4dgG2owJnk6NBdrIzOU3FYMQTXmMg0tqzdRFUyX1eVy4ONdymVgmrFGu0DOhYWrIN8BuGNk88ihdaSf5daCJoxJPpTo9pqKc6iBVOpbPKTxvUzrXpnZtsvd/KMGQaD8LykuoYGlutvHHhks0gOdjHgz5bGp7Q0jHtFdhbD7uItWUpBM4ItxT0KvWeh589IUtcj2T5Lxx8J6AlF1TYgiQH9jA8JdBAWa0mAZMFIYMIWmgU5ZzW0Y2rn1JKBWYfEblCTlY2W0V6ukiuBQBmtGWNPRZK6xS5YzzbN8GqecVR37Se183N14hPAJxaHSt77yws6DtW0pR6J3bcPbg3gmyM1mYQeJL41e+O84OO7lsvwTFS34VgtU1m4MwnXq6r53qU9WOPpi2aFBZEiexTxDjWq1Rdu8p54VrRKG7V7sXkRRNH71bLArVi/sb7+2vqC4/qXYRSjg0hHy4rWKeerWB2mqslqmqRgo9G+m6kg4+2CrlXXzjV3UXCqopZ5jgUOV+H7MPoQKuyVWA0SnlxY2by0i+yxC5pEj9adudu2EcXcZ0+3d52M5n16QFGv7tRsrEiN9V1bbWubVfeQkvbalDkicnHHfzmtaU5FFJA2ieMortd+ZrzIY3+NzyPGnV18FwAr4F3kh8L3UDaYPoJGfrMks8Iq7QUYf+7Nz2HqB39jO8LrHsHYH92e3LRcCiZkJ81fbP+4n978OE1IfE337vBtsCROCEUTWfU6mLixT5JsH7NoLp5+Zp+3fQnxNZYO+mcXJx++VfuFG7/HbfknlsRu+7cViW/egDk7S6P4URDUa/rWblnyrG51eVVCeEIkwP6iFaZqEIC0Y7KIrkm9YdLkgoDanp9AHUK0PPSmbFq3m2ynMNQiSR+F4DYgf9/G7oJIW8DNtCP2Q248YceYt3VPV37gPeKbzL/1L1ex1vIzGioRdd6mJyp/F7tSV1mU9fKO/AnLaR6xCDZGefBL3upKt9Q/dhMy6OVAUmIB9imLEymgNE2GpJPWC5iXdMpahoyD5xMwYuDRPegcXkqUYdei70qgeZoMeWOAvDFCJgFMAJ4BXMuQcdQQWo4y07aBSCauv4Zh8BWJcQ9IjqAka5JckW+pgY2LD88k61eRahnQZC8fpLAT0/bK91rJTQpT5ZydMZCmSQrARhT9xIOkJcUtFBphP9xKlh1JKCUqMNMrdExwVP2GTQ6rMFktl1Gcgg3k0Wm/UdyrTvXuAi0ltVB2SETTyoYstExyUYAeLt8HjwmzVRzDWK0mJmRJPRhbdmG0OFK2UTdX2baIPNzoSQ1tgxkWwCZNXn4e+M/4QVYPxHe2rZ/yryRu9HqqNU6j1A2eREGi1TsKfsFjVLSeznmewaoDUi34cJoACvU2bqJLAABwcmApbsaDmLgdF3LaYDqBQUV/UvtKgOUcMZHldolct/z3QyB5QEUMZeE2/1A5Q5HVnFHL8BSFyyT3OvqgSQ4093viX16lBdHd7Cq6m7uI7uYjRHdTLTpeufz3FtHlVZdkh4iNElXkU84Mg2pvcAT+kZq2WOjtplw+2mC9g5A0DElSxcLPqJxCMD7OpQNAezrXMxfctkRet6t70Wy1gI7XnsXETck3AcGveo2B1rIzVPSzTY8p4sbvXDUPrA6ejRBbDxXwKyrZDJ62hxleEuoTiouLpGSdKrxyqnwJFXLrtY5HWTQMEny7NT2Y8gPfla0EifOAPRifOL1hE6qmkbQEImbO1/nudzneXLVuZdycZ0ThPKu2K9/kDwUW4BfSWQBpZ39hazjf2q9vCy8sLrMNMfT4x6NpYioxy1R3JWSI7roCkWYWGDY31AmTb3Eho9AMfDc73W+MAsH99rhzWxeOefGeG5B09/sLQbLibFP+x8R3E0tt5iJs5kJpWnbb6Tcmu1ZH4QkTD/D03mF+KG9rJJwvluEJIYPq+KGhCe+iZ+zk0UI+pbOtUna7v0OLVLRwC6tDWxg5ZV9bIjB7hqMdYLw/049PSjFdOgCDR6mNwMphi4zEwYnsIcgD8cPqgXqic8hdiHw+kcwR1EA2NtnahoqbHP6mGp4PY/nUb+Gh64DA6BzjiaVbafGEWwUMNkstBVd+Sg4C4FPD+SyT13lR6USU8USxhvkAYP2FLt5YR9YXX+TZJfQM29jl+IE2t9xhe/teYdFO0VFFB80/mTHlXbvhjDyJVmEq6959jSpuEQndAvvmQGrczMDB1NzOoWDCIFJhS7VaMdhU9c1HBJmP0xPJUMO14sIirPapyqZYDbWamuxU3ltbeYeuqXCn00KZAZV6idAa1fglq87a5xTsrffl68YbbbjVdOckE4dxM4JWWqEsed6w9opbNpRxYF2utvdxoqj/U6GyOEQdnOQ+Tam+lmrrulpbUZxrVVfXVbqqfKy36Om6SktLdXRt1NF1uY6hkFBDzVJqVCEb1XNndVHM17WulOsqpdwrLaFsslb/MQzh7cAPyS/cKXEmFYBJGkfvSclKfhnlJ+4SgZPfVm5MtkL/W0RNrdoCV3drf+gUu1c5sfHKmjacsFX07JxYC02ypmH02waT5T6UddA5t/7yFzRTcWNCBYY0rpahlE2ppnnUKZ1Hnc/z6Od5tGIePf108+jebpOnUzJ5Op8nzz/15Hn6CSZP+rccDMuvUyGpEqeoh+SD+K1uVpIycEI3BTkafB1bsFAWacsp3TFKhvtNtkfG5MhdYY1WjwhmkmD7El7hshZednBTdsaB+snSdUQ0+EADMGIiwys+nInheL2MdVOCpQoENznj3TWIz39i6qnckfMYf36J0jlA7qnD7akyRUnRbQnLcBWOHMustsLEqblkFaRSzPuTRVbE4jAjksf/oYPuFE6xOG9qSJzSm9wxfKMgbdQoCmo75vPLpgoXKfhq/CaXPo7zp8xZarWKFa9aXMpgKLPW9mGeB/UibxWsEhWeXv2U4dAvGU+vX/75laA3KZx1jWJPu8NLRjzkfBf26OeXkCGBA8Oaj+gomM/7B0M6zReL1OHSoNby+PvHqwj0PFmE8GlUkZv7q8jNTiqy1V7VdURGqFSSYg3vpyQK4p9HScQla3p8tGkZw5xNzk4etDyvXKDML8zEeL50SVC+AQtDl6H3BBxOr3QZ0POvaw31IkU/pFSldTl9RYUVfHf6lLKCxlb8xAJfTpre3pmtRtaka1NrEyNothK5AyxeHPwtvTeYevPi5uAS6DR2wwR3mv8Y+5csAJBGSxgD5mX0YeJ6FUdLEqc39Zq/cC9JKyao/X54ife70D03ICSv1tiBQKslbh9t/R5FCyTg7IiIV4216DXJCVgmFLO3XNeEuFlzqKJ+O3ODWf3ajetauWg1P7iVl4k3y7U4EihTylpiN1IMvISWuLyWniJAmYF5UBPKwvB3biQZXJETvUUZhWNitVZcfNaFzm8Afg5lZbV+cFsY+Df09sZ2nyxwnIX6igqX0PspWkrk1juRK+xiCILn7pQE8uYOvqpE019p5x743ILr7wvoFxSGhtTo0nmAn8r6OU0RGoQ3X73BjoTSX0rxNBmK3aP8mN7FjXB224EaSDcqG7FQh9m9R9JQVKOiwTvbHLosW7fbAxxci43YQI1jGWqXaIizKJs/fs8QrY92pLp61DRJQu4R4kJno8gCWR3LOyH3WqT+p7cXbylGyLHtr0DfHtz6qH9U/cx4fDyRaruNE910zSOUBuYwCvqEDfLCiaAZJlYQ9hHez44yk0YEy6joXMszj7YwZan5WQ9ju4vkzM0fvpfqv1illFEzjlIgj8Nma2x75LJmpG0Yl7lGxTjsG4uJlfljZ40zaWrLur/m31mDi4Z1pQoblFOVyQ7auada0KCqz++tEjn61t4lge6gQgXoaT76lwMJLaiCgbFfrA3VsL3LIRXVzZSyFVGtRCVWtFJTZbn355TvO0QYKEiTMYO5pUBHUj2aeMwKkyjtTaNsByANvDyO1iQpt+A/ykfIC2jPAjdJnvtJ2gaLBSzdcB6hKPkjDZnldJfwEAdjVxxWbhPlXhDbVfyGM6UpvraSILG+S7Xx0ZKadIWzVhanUX/7n+GD27yPbM7eaht48HmMAmtlhdJ3NaReyZB5FNWqsac2alqu2IpTlBIHYHvEyrPdtSpuRwNIUrKU1zbMomDiZCjbpFY7t2qanJje3FNODDmTE301paZllouJA5SJSWSXiokD6GLKkoseFtqrPXt2hVerJHyCyGeIKvkyokXZzTBwCjlPYLzBZ0zoWxXkOlUM9oIsAKANVb0kKUvKQx+abpUD7qkr1eXblvn9EjolpXtLcaGS+L/e5DD2fIO3NuBARGC4FLrRNAlF71f3QS4OuhjFCYDCJxxyOcV8/BEpbbqRrpbZHwt2HK5WOLiO93zDUIun+PFGm3afHqh38AYaeYVArcGPSz6NK/xnZe/CekRJyAOnSjzTqOzScPm8DJpTxYFgx2EgRzYOBFI26+k1x5BDO3mtY8ji3bvmhjfFTFErw4EqmQRJH6Vp7E9hMq/XaAs35abVAnHzCN/h+nQBPk7ROHNrINVOvYASo/l/RPi0VksnwrNzAZjzawfUn6fHkudRoxjHyB5xKu1gRSNsJ4EIwkwiL2GgwnqK1NpE50GbyD8dDzgYFXnAy7A0oLuF64ruh0Zo52hdkRK3wLInnj5hHFrQLDMxMwA1PCfdsb5+QV99g4lDG+kwAIQbp2isH9sL5qHH+OCXH14+CXxg5rVyJlha8gnBLXotxPXgVlDirkorI01jLyCot8YVIREDkGz+kp0gmYkrwnQSjrrEw0Ha1EfQHKiMawO8HGs1u0gobfmYNtaC2rlchhqpLcJkSzgZvtC+01x0LEW7eWsbu5qdr3LPAhV1GqRoWDQLg+/1lt3uo/3ltDtkod0Q8TE1tFQ94fHbgrLkQuDaggenZlfyJXuGvSj3Jm2rgRN9V4mi5/NVICovaFdpcu7fbtNjTjRDwO08W+UpcaOempqnR/mzHQKKVrkpCqNfymscabQ0YUFyjgQfCg6V61F+ejbDohk5Hv1UMMWbk0VUlpPjsm8JeWN4nKVq2UHSAxr6bVmKSOThqGrBoRDuYs9l0ghcVgAurUj04VOJb5XQl4ZOVYLy0Bnz1aBtxLIVBJkaD+JI5FhKTm+jnUrHiQmooflCtwwl+b4WprWrJUwzBPO/jaPFG26Zmo8RKrsD8KoT79WaXZ7B18Npav2DH3rRh7ZHrsHDoKVSIOVm/hIYfOpK2dyjF+NYh3JR+LmtOPXxn2UG8Mh7t0rSBYsBagWVUi1/bmeWJK/WP1G3j75QEifsPH9dt6EVOchYWt1Vgo5eHcRW0B9qlWhYX0nvaZwUrjf5uCVd8/2Xd6LprtKoVtJMdM0S9VCuobwJ7aGxJcs61ZYFZ+hjtEDp1irZwlKW0Bh/H9x0diXprd5j+KYU4vnqLcNcrRaI/gJz62/rMaEBZfqG6YPbEsXbeEscgrSwSMDd+kL5Fiub3xmjhQBmV9RpbGbo0mxkGg/k2UqruBS3ED9YwcXAQ1mpm4ITWR60UCtZMbhlt8cFHltnlsizq5wgWYBW13lPablkFkdB8FixyW7ppGsqAd8ypBw0rSm5cq99fMW1hvcruWFa2+hlyI40LedZmEZ/88mH+q0BHWgG0ew9pITEBSXKCW7UQ+9FeS6iVUI1A2Wqx9DE1Ue4zxQkJ284xdAWq9y/N63841dlTBOYhhPQmTvKuiV/ypcGfqOQSNdHal5cGWDZTdDZaciLa+2OxGzzUVazXQqlgi5nj61FK1b6FneTbpqlB6JL1tAKZaSa9bK9CHrgY3sR2Bzqgpj52kml7HIBV1yZ05jsoAi6qKXgBJhGADHQ89BiIp5YKMsoM9VkyQcZATPy9yJKoGHzdBM6vVAs7w7Q02U+Ti3b+vprRWQy8IEGzFwviWPVQdSkxbc+1JTLbjUY7h1rnuNuS8dqtUTVNd0xOG9GNo0c7FT0QVnRZZVVJLIpaalfs5b6XjjnFU31a9ZUGbTcVt8bvHmNN9pvq9tKrBd/ZGP9+jGNVRxe7tBWv96hrfLF8VrpEZCdpq+IWQT67LXjFLPTBLOVERqn2zaHvhZHDe45kwr84nRaPBWDAPefHC/yWZEVOtGvU724Ll7/Ws2F+WaPLACqzjy1J5hBvCMMlmvauOUyZL0yhetOTH92m6ws9bks85Ei9UTVNpmUHMQqoOmPAphvCNlJVz9cERKYdFWMj26QQrnFwJZHgtT9FV0Mhttos5ScK8Tmd8s/JXMXFEcPieYPUReM9gY+vJDnM9raozeVlr4skyIg7mvX2OWuAcwaCPcc+A52dNmz++lkRIXQRD8IXeTo1Optq13rxOrp+qeUelzCLt5mU7jORmVX+TxQTzCajvqV1uPYam2tyMG2ipyWVcQP71SR1vaKGBYzZRK7+8DFH/dbeik6lbd6JJd2zBnF/VWO+2EI1zrUdPihNubVNQkp0EoclwWbpcL+XV/u+ZSlmXzZvL0Mt9Uu3STxr8kRe71lo6yJmXZRVDWgMYBRuLCWBap+DMlT/mbPp7uuVr1dSL5o1LD5lQUvOJBIMcBl4WIZlCVO1OuZilfo7X6JXuEaPfUSvOK1eYV8xkKq7xYrvxiveC8SfU2CPdrAJ3+8awK0BLdEHEmvLIHm8FlXOxmbHSy6xy7S0iM8zsxd1sqg8tM7lWALGnFkbV6z28M+WZTCSharH+LtHS3Vcy6sK5iPmpiB0TieB/QwcI1d/1wOW33ayCCx0nMUOiTWKvPMa512v7ZTfHdb+LqszqsUC6SNtFxrRzb2qs6GZYe8lE5QqR8ySJlyyDAlG5D3tEXnn3BaKFxijjnfBAVffEeNL6u1ICsbRzxJr3tH3rtcADMefrEwKI9dAFzNbDnTofu8YCYSlRU3DWzI4m0p/crN2wVoEVxp9at4zr1Yx6mCU7a/z0goX6xTLFrRd8cEqLpL+oSw0fWBzQB/HpVg9f2XaQWP9fzLVUJM/KpWlG1ZytqyeIc/P+N7y98K0Gyio2ISs+MM1+kfmRIRejPZ2zSoFZZe+dQEeB1F6cvII/VG+ypKUvA252HSFgEocSMj/JwcH4Ll7C/T0+PDNCYkmcEU0IpXYeuKxOT0GDe4W3SH1cn+PAo8fMrjIgTK+6fHVE6nx3RhyULL4WR/dkVm76EO+0acizS6vAwQ9ZAiqeTpbR8X7nQKXvB+VdbXQTp5567bVBAWvxa+3muiT/jPf/yvlt22R9b/+792e2Cd/fMf/we+x5D5z3/8b6ftnFthFP5O4ujIGUDVoYzT7B9eH7lgcPjDC7KGFI94GlN4ohtUxruYoQpAM1bmQybu0+R5QvwXmeJoyNorGPunr8VczJqu3W4Lvo3NRtuUNwooauDPqNIeRrOUpK0EcNzF/mn2shVTfKo97AuVZ5IFL+b0Qv9K5ZowsDbw8CaKQvH+BSeRLnl4AMdF5dGPei0liyUaFkiHxDEIDEw4PBqSG5vs/YJ/e/PjyzZ1detIsM1PcOv0WN1rDaVXAzVEyd7/sNQe0Ta9XsHYzZ++wF69pbuIuohmVKq0X9Eu75II9Od2H3cv7h9Z+09ZgCy7AgfsewzrgUdruXSpvtG2vkcT85Cl46Yr9m4BDQnuN6196bUCpPjdD+OX10PnJo57Tid8cvjry2/ev/ohJlc/DX4+eDY67P/Qe/6qc/PmetiPF70Xh8tHH1bXveXyu+fDd3/tvB8v0mv/Ohi53w0Pwr++7B08Wj5/dX15cpIVRR9UwJJ4j8QM7fECzH30128eF//P6Ei3JwDo7b7rwD8OFmDDjx54I/v56VlIObvdx4MDSBZAIZd2YfoJUJYDOJhKTw1AKrg7+9CNKKlz+H1TRssu0rK30NKOKyFBTFavMKKVEmLKLztCujARQ2ILhqhObzjs9u1Otzce9bt9zHPXKAdTFg+kQjbz4Paza6eRx7bt2HZ/PBj1+k5vPB53xxyC3kZdRlM8y4x1OBsP2IGALn0v1un22GfPoZ+DIeQM8P/0cwQ/h/jyLfscjwDYxif62LO6Nn4jD47DXtF1MGGI70D2WUIHEjp2FxLG7L3aHh5NwBI6XQZhd/BwAv2LFTJA9DE+G2jbNMHBIwxDm6ZiwrCPLGNJI4rRxZ+dbvaNxHCC4OBY4gArPaAv4naBWg8pOs75OSqgsnkfRDRkal4WA2eaJdqcP/q3L3CoemHn2ccYTSXYmQpC7xOhfQ9ULlfa/Bkv1EyRIBR1Wxnwe0shzscXcr69rv8ZWjDeq3Asc3O+kYdgPtzCT5hcT/UJnf2zo42hTfV/t57S7UNH1pNXP1s2J/b1ZToxl8IYwCmeWp0w9LPgwhEaHGTf8j00huKLfEaYjYfudDgm3c641593x87cGY5tezgeEzLqC2sD/kZ6p+zZHrAXFuhNXCjcHoVRWm8rptbRlZvUT5lF0zYZYRTniNprxMPLDxp0sg4wpPilS//AhAe1o6X/MXagxYu/pz1olIWpVQE8cJcJZqLtaFIMjcaXFgf6JHZmJUPlGvUH26O0WvWzM+B7ZHd7/c4YhjL46PbsoTN2uvRj3LNHnfHAalrSbJK/83oGqZ2B0xv0ejCK5tMJg+/3+oPusMvJDnvOoGdZMrLTtjsjZ2QjPIAMhuPRyO6xDxiyRz3b4SX3B+PeuC+9MIts93pde+QMRgzZ7o+GDs5NAN+DSoyGFqPUs/s9xx7hg+3ci6NKyhu0USL9j+vdMC2NOx2Q4wBqMbdHo2HXG85BOp3h0PX6dt67/2c6WF2Han5/JFR/1P/sYn12sf7cLtbLn5KD69G77395fjhe//DoYPTDDz9dxzfpYdwb//b7D4fD9NGz69nT9dXvh09fOL8d9tfO316F4JbFg3D4U+/w9c31k1fJL9+/fDc8fHeVHKbRd9evPrtYn9zFsgcde9CHwbgPzoIzGkkuliFrm4vVhbmm0x3C7DMY290+ekiaj2Ug+tnH+uxjffax/mgfi8wdD+zD7nzesXv9rgOO1nBOevPReNYd9KfOH2uF/c93b3a18v7bOzhD8BTG9pj5A6MeOho95ll0nZ5td3hOZzTq9jt9S3MzxuC6jDs95maMh+NOn/sofOwfUdeHj/Ia8qgPo2qnywqD0dnp0AmAOkSgmn2L5Qx7/ZHTHVmqa9WFsrqDXpeCOMNRJ6M0sIeDvj3kdQCvYjQAtj/WwfloO5fd2fLMoxdG5qbZuDv0nG5/ahOYBGfzgQv+mDvyiD2Yd7pTp0sNk50DIzLwVj9LBt46HJxPuDkezWKwwI0GubCZUVYciINrVvJ+OxPnBWbvF1CfBBHdLCG+2zNMoMXQt18meU5MYKCbkV/89KquoGdE57F7yU9iaivQT/nntxwC6yCglVVMQRjzDYe4s8bVDnGzV2XFKcuMMnzyte/HN8+AuMDG5xUm7D3iaBXPyFN6fUmJDL/E8XzfOrA0dH45kCyWjFp77seJKJvWDBDy3NxB2agtYRTyNn/k/wNfwAcs</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_277cbde407034bf9addfa8e4969bce7e\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
+    "    WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\n",
+    "\n",
+    "\n",
     "## Restore when checkpoint structures differ\n",
     "\n",
-    "The ability to load a checkpoint as a pure nested dictionary can come in handy when you want to load some outdated checkpoints that no longer matches with your current model code. Check out this simple example below.\n",
+    "The ability to load a checkpoint as a pure nested dictionary can come in handy when you want to load some outdated checkpoints that no longer match with your current model code. Check out this simple example below.\n",
     "\n",
-    "This pattern also works if you saved the checkpoint as `nnx.State` instead of pure dictionary. Check out the [checkpoint surgery guide](https://flax.readthedocs.io/en/latest/guides/surgery.html#checkpoint-surgery) for a code example. The only difference is you need to re-process your raw diciontary a bit before calling `restore_from_pure_dict`."
+    "This pattern also works if you save the checkpoint as an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) instead of a pure dictionary. Check out the [Checkpoint surgery section](https://flax.readthedocs.io/en/latest/guides/surgery.html#checkpoint-surgery) of the [Model Surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) guide for an example with code. The only difference is you need to reprocess your raw dictionary a bit before calling [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179)."
    ]
   },
   {
@@ -321,17 +371,17 @@
     "    x = self.linear1(x)\n",
     "    return self.linear2(x)\n",
     "\n",
-    "# Accomodate your old checkpoint to the new code...\n",
+    "# Accommodate your old checkpoint to the new code.\n",
     "restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\n",
     "restored_pure_dict['linear1']['bias'] = jnp.zeros((4,))\n",
     "restored_pure_dict['linear2']['bias'] = jnp.zeros((4,))\n",
     "\n",
-    "# Same restore code as above\n",
+    "# Same restore code as above.\n",
     "abstract_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\n",
     "graphdef, abstract_state = nnx.split(abstract_model)\n",
     "abstract_state.replace_by_pure_dict(restored_pure_dict)\n",
     "model = nnx.merge(graphdef, abstract_state)\n",
-    "assert model(x).shape == (3, 4)  # the new model works!\n",
+    "assert model(x).shape == (3, 4)  # The new model works!\n",
     "\n",
     "nnx.display(model.linear1)"
    ]
@@ -340,9 +390,22 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "    WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\n",
+    "\n",
+    "\n",
+    "\n",
+    "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_9572198988a54a0b9731fab9e0c97e51\" ></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_9572198988a54a0b9731fab9e0c97e51\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #cccccc; font-family: monospace;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtPA1X2sq2f2UOXesIVwmET8HqeqCAttVWsbXtu3dxJ8kkGQmTmAwg3uV/f3sm4SMQUE/V0/ZdXArM7Jn9vWfvPWnfBnzikAOF+4QEuuuRnu+6HP0HeW5AOXVZHfnEwZyOyB4yXcazJh5QZ1JHA5e5gYd1GB/blJOs/FJHng8jDg14Vm6d5RMPRpnLYFjDet/y3SEzsrrruH49XLqHom+aAwCwHzW4XUcm5QDGOGF8D3nYMCizsg4xeR0VdFsgYSRrE2rZMKIqZbEN45gCzbNl0YfsiAZUow7lQDkecncGm6WM+5QFVM8G9I6EsxG5929zoXjezsST9YcMcPowFug+9TgS/O1vYc9zqI6FxHKuzong3id4sHWQTmf2D0CggC/gyCAmC9A+4jYNFIvwC5D2mWuQdEax3YArch5YIxz1PMIEyw1d7CoW/e+/kmaOMTMcAtNs6Dh7IQYFyOy6LoPR9Nj1+xm0SIN7BUNiKjbMqS4GPeKbrj/ATCcKc8fpjNQvIEivzKBsuOgtKhYysA81UXqJasUhzOI22t9HeQGykXSf8KHPQO6IOAGZE2YPmaBseevApiYX9EkA8eEeftZgSINVMcMdKz65GZKANxgdSHW1fTwg6VAmGbHH3goibxjYoRj3EnicotgP2djA5eNpEFSEiuSuZTmhV/ak54C1emIvMUIcvoPICAw80qSgTn5X+mQihJ7yU4KgCFjRHRwEH8A5o33TqdmevQGYYWqK/D4D8gTzlzZ+8DaX5AAGHSG54X4qHj5SiGMNOCW3+6l8ClzX56sgLgMSQRgMpjY5Q7IE0mLNlPcUOGMYxjSXc3cgAkOduTytmK5jYA1WM9i2buMgfeBgjTgH8ZleiEOuqes20fvEyGTQP4TopoGHu14d5RW1TAaroUeM3WUly/B9LzF63isyyvWwpvlkJK1bBr03ld0CzufnALo7GABbCxBYvoRqlkBwSLLtjoifSYCPg/fGNmE9cuuBQRLj2SS0DivoF5zGWIDIw8s05xBsONCIvwhQ21XLlRkAMXpC0nOA6IBYPUdG2E9ns5rj6v1wCAxZA+TEDxWnercocB1qbIIMrech4Hthl8QXxHlADrFBTjEmcFX8RMelNMw6ohzD+SAWx8S65qAFPSdJH6ANGgDSyfRAXQZEB0hqr17XCMTqBcm90eUr2TTDQzWrilM1OnzBHme4KJP2LsWwAae0wlXMBvb7AcEW+D9bXf1MVjijQSxNXjSFj1Eos4M62vpnoazpW38nefFFa4msvAKRQo8C8dAPhAI9F3Ik4ifgpcHzoZWuIBFlZVQP1tn482Cds8fJLV/FotCgZ1I/4D2XhUFo1bU2uZJSKAtvSlQV+mHyQ40vkyi4GmDfgqw2JEM69P0PYoNQ7E20IYRGlhiA5tNJRptCqSUoECSUCsnA/yRqyUgtVRpbpxisgmIHdScDzXUC9HHIBb8GOgxXwrs3AcfIjonWh6ohjLwDyBhsOLohp2ccllMcEGNWa7whefGzt2rm4WpZDOSVmjja41yG/pHARXK4m69Uxjjo6XAOgGBn67HJY6fHNE5vwrm0Jo5yUfTRAWZgjrOYgWJluplZHBZIRA7tYza1ZrktUgNEQGKQ1WTdIX8aKzMKQDGUGH/EKZEo0R904Lk+x2xlb813+5CoiJF5MHpYugvLFuQ5VfO9Yut2TxapPeEdC6ld6Cx5JfTYBThfJHgLgH6Y8M0gRcIHrBo9HSokwycsYj5egwKVccBny7xmzhgdzTHn17Gjp6FQhioNEhqZ2CsBx2L9jN4XoyTKqEJKDMirwExEWrWgjpshdhjUPT3PJya9hU1ijrcrHQ+KPCwyrTH2GWigNz0qpto1TayrxQRAD6qkFcXJeBgJKRrKRsqcp5R1md1jP2v52KCgtrRaLBvE2kEu+IhFUB6oq+j2TugzUN2ICCSHUCTlFVJWQvWzCT6m8uzMMpcVvc4wl+FmeXeoQHDtOMirm25eKel25jG0vo4th3WBpCqxOoiVi5IP0Iljo3+gtArsLC6J9g0AimRE2R/JYL40+/i1U7zTGkaunG05Hc1Gw68ozXkVJiWYFU41DCIGVTIA7tYzlpBAPTl9ecjE/xMrFEP9ilD1iEJxM/C9Aju+SC9iruQfaUfIbgAQ62APEqOHi8uni349hrCklDhiTYk1MM9BRxKKiNE4QM/1F1ptSWJZD/48qt5BSw15Jd6mQ48keFnKM1X0YB2QwH942avy+0jyY0rdrMjHYH2UiB4jy8cge4ip57iREE1i9EfD9/FEMX13kDZcfSj6g4pIwAJlhJ0hCdKZjBK4A5KWaZloL4t3JSzKRGv5kWVZagtlUGbWzA9sQrjo+JMxOux2u4KbrhgT/Xs5qfhE9tW6E6an//0/USmok2mC+PSycLERx8QNhhONjaOQWRLt18DX62joO2lRo9TFfG7smmZhT4MqqFLaMfK1zqnVaDbk6+S80XDlp+bFGP4etxuNVmPTqzloNKy++944aTUPx98ajctvh+8apyfNw0bbuj05/mDzoHlKiVVsH30tfDipfBt1vSH9dFq+VN99Pbn4cjq6Or3jnybt9uH2ldW/pM2jvE2PzofvWkbnOn+s5czRieHdvK/YN1eUng9PWcc+Nj/zxudK88wvNdonrN+q6J+HQ7Z9Ub7Rg/54ZLad3M2t1XJ3Le3duLOrHjdyrHFR/uD779SLbesuf2HkG+9M1TqrHo471wUr706GF9XqoKVWxsdfax8tyyOX/UmJnGh3ZV3zP3Y4bljnJ2fjIxxMgvPhycnXq1Z73Ph07p18Mz7ncttW9bL6tcjz5vtPN41RGfb80DirNk7HjYF1d9HdHn7vktbX24JZ0e/OShfHk/Kw2Xh/17z22l6RHp8ftvLfh59K3Sozmx9ax+3TQYNu745aBZupdnVb+zL+ej0+9kdHnc+H7NpstSy+/VH/7jjVcu3w3bi5a9dKp6edbrHzvWENTsrXzfMav+yQ41qr2TzpFI+s0kXumz7RGh3Q6Zf3ucZ5BzfI6aHTOL5rfbS+c6vS/GR9/Hhy1OzT8zJpN78eNts6zXu273oMbMP73jpS79R+1zw0uT15z44N3A6OzfzZoNM6qzSNxs2XLx7mQff7wDAwrRXMu1rpM72+qXgDv/LR/XbYpX5nMHrXKXavusV2q6A3z83L7WPH9TqldjAuY+umsku/k+6Z412x5vEJMU59Mry66RwO1Ku23+92b8uFytVVMG4ARRkkb+94ekua9ZbIUP4Nf2bejw3Xgzpt7pLyzlFRlA0QO6HP/gv22nyLY8tLMFmch30D2BvMg+koHZbv8StKcMFLV7gvgEXlvRgLIDyILUS/QhT5eIwpRwyPqIW56yuws6e52DeUsU85uSS3PD3fSyRw4V7zezDIqNKphWaGuAEDLJd0QNwhT0+vSFfW+WTgjsjK0vsdVMjn8zJxheALRUJatuWS8S50LFJz4kRDchrBxKVhCr1BbUwdCGzcRQL4DxnZILlnUENDNKYgM4IN0XDZXpRddJv3wD2eaN4gGR73U7Ecto7cvqPbUMFUd8si2cwXUUGFb7VSpVqtFkvFahnlIOEEfpPyd3GLkIr2jy4K4zcHy1UxAId5w1vKvGF0kqVkpqC5t6nETaKkAibDhAKYlIvjeGPXSVFajuINoSVK42d/6uBPhwtBAsRmuNjkUo8jdWA6+FZhTPwqYZtBmW4avn2Qg+nZYCSPxU1Xavt4QZPaAL0CqlEc7K/X0I8oJFkRLyxWqMOooCRYEusn7ONBOkmcW/MO4hZy2aHw5f2tJwYzeR2d2UKz9uZ+ShGiTc39ajYFmYucE/3q5c4szEnHhYPBhs8R2wdxXmKiWGlBpjbNpw5kcrf/CFFTYz81u+jEpYKeL6g1raqqpQoxNII1Vdd28yU1XyMVdQlp0gVp6uAadCSzTmQ6LubFQrq0k1libZ0lxPLiRBIrGKJRzSxUqgUTwlOlljerBVKulEmlWtvV1U2B6HnMPMlXkyUhjD5RGn9ac294nGCmNWDqYOcRMI92sRf1irDC2OQbIcTzeEjsEYjUAYLT9AsNIILQO7kjsqkBPgKHKJpXXqIk+0s2+VuIb8NbZnM0X7HfnVX4WRkPyniizf6CwXq9ET18Mvcoo/xVj+cNzwulNmrpzcOqjLWpHjrRr+WBrggBiKvbO+IvH+kw5Aab8qRN1D/12EwQUvgcVxjMzSGTj/ohSRTCHOVv1VK1opHabn7RKX4GO5dWtcnYJcCLBd8jULBo0oCQZBZeriDXRLnPAag4R0eTOxCLlSNsFOREWgdFCjZyDtVy3oTbLisqqpoLxAPHkLn3sUWCHNhKrhf4eo6xXMxgvMnjA9svFKheWoF/NWQZcCBahBEfO79y0EqMRaLAiIAX2PzZws8CadMgpJJ8vlj5qYLQApXrrHgB5NUCUbVYerZIBODi93eMQK+lvb8chWRASVB+9Dxy6uDMZeQXk7ngaa20xeTry5mynkkwH/okSJJ2+Gx36qD0a0l6gat18l4AeX2p94nPiPPflt2LKD8U7jq9h7O/UtuuUs5Xy2apquklXCpUq7hUKRarul6plXS8W8B/uW2HSs/XuMvXSJmomlbTsF4ySHW3oqmqUa6ZRMXFglr9FRp3II/fv3UX2v/m7tMizH/bd3+jCP//tfBeNnj/2Hmd0MjbUMXJE5HppCduUSmzFDHruPAlENqRTbGovKtWdaNa0xfLu19KXRtbGQsgr1YEFtXqq7ejHm9N7pD/jkn3IlvrjGER5vXd2BPJaC+htPwNWltgNN5kmnxHOc1P0tiSE+jPN7eF6t6c1IjGcPhnamwtWMk6M14AeZHD/ddy/L9HYE9we5/oNAA0v1k/acbXWqlPAV5f5sOA9FYfS1oS+aU/JC9QWL2YwKdMrZP3dP6ZY8LaR+Dm/1EGWvrnjqkn1wePr3NnkAYdHfwf8V6d0Q==</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_9572198988a54a0b9731fab9e0c97e51\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
+    "\n",
+    "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtfet62ziy4H8/BduTbkqxJIuyLpZle75c+pIznXQmmZ45vT76HIqEbCYUqSEpR26P/u+8x+4D7Cvso8yTbBUuJACClOykz+z3dTI9iQhUFQpVBaBQuJ36wY2VZrchOdv3g3QZurcnVhRHZN8K/LP9eZxc+mROkoT4l0P3aDQYz3vDUW/eH46G4+581COD4YAMR+Njz9k/P02XbgR/I73zjpsk7u1N8OulF0eZG0Qkse6sj9dBRtoA5xEsKFm44cTaWCbgThDNY0CZQ0p77i6CEHhbxFFMsSeWF4dxcmL9waV/JtbCTa6CqD2LsyxenFjdTm9AFhO1xGVC6osLouUqu8hulyCQxI2uyP4UWLghSRZ4bth2w+AqAi4C3w+B0jwIMwI8XAG1FPJJw2laMRQVZLeNbmfQvHdhJ9fxDRVUmfT96EWrxYwkQDCKs8bJPPZWaRPIzuLEJ0k7cf1glZ5YR8v1p5FkvynTSF7oZET/THhxJ5azXFtpHAZ+kVVTaicFSJKkur3UaY+ykAVLwFEMeWIt4zTIghjU5s6Ah1UGaTPX+3CVxKvIb3OWaUEmhmchwAIV1/eD6IrZlXeNZIMINNQmNyTKUlHYx8DPrk9Ae1kbmYOsiYWczcP444l1E6TBDA2nXK1f20HkkzWU3O1262s5i9c71jJet9Nr18eiu/R/WC1aoRZP6EECr7q5Qjlf4xq2vDDwPvhu5t5HY2HsokQvFyRN3SsiWY9o0ZvTQ9aXnGYJIakXL0k7WUXta5JAWuolwTKzqG3a7nIJPLgogcPYy0jWTgHHXdjne/gHik0zS3BhnVmNRtM6O7fu9iwL/puvIg9RLZ+kJAmglf9KfgZpHDew0QCAZSUkWyWRRVOfIJ3OPIkXDTeLZwDUshoLSnDR8WKfvEZRPska3WZzAtibvepivgMxZEe9oiDG6uw2Iynw+aDyBJE50kYqEflo8YIorQYl35mt5tC3cxReQYazjesX0X8Lz7SYe3HMWA5JZj1DY1q4yzffP30OljnRa3NFsmdgjEG0ilcpBW7cuOGKtJgZAiaiiRoixZmbkkvaGlpWPJ+nJGN8BHOLoVqnZ1ZXYFgSPFSnO+GpDLNI2VgkTIlE5PzMciqIyJx1QhJdZddW2+qVSDsdlbggxkTsZWlOkRX52GqYSTvNiYmPl2523QG5g8xyYs0SF0U5X1sO50fSdKJV6KIoYnrRnSJTDrDAyDWtA07eqkKyDiyHI8raYYVd1RXmPLQwx1zYrK6w3kML6+mFcfu/SFrWVcuaTc2N9jYC18l7knhpEF2/IUC9wcv7QG5pn/9XbvZhsHzpgm0n7seXQcT+xW9O4nt3Kcwyp55mLoxnb9E/8XkRDcDI3FVhwWjZXwXpd0EEA0ODZv3jH8yEYKhqrJvWISJYp5ZD2v0CL6/gWliWZs05AKWVgmeGxB5TYo9zGPxDAcL4qlEu9YBj/z3JQCn8axl/bKwZQMvqNZu5bW8kK87bviJH60zpAzCfyVPL0MSPFsPkX1lZHZ6RNXGmQ+a1VjJEL0q1LKAW7roh9M4Zak4qanqaQ+Rc/pv+Yc0OtIVtxjrb0zRP1stGbgKqDKAZFuYNtpODiarlqs/LAFGJkg5ZQ1bFcqgJjg0eIFjoawHX1GTypoetYyKZjsDSTCcnBi3G6aomkHev2NSXBHtZtHYYd+sa6yGnKWjxliWIGEeTT9ZWYfRCksWnyVR1QihXn0TxIojAx0hyGw6ihmQCpmprXR8XAWVB6u1aW6iIhoLoitoUpjTdqQwbFciFv5PaJHLFEFsaBbL4LTrSb7MEPG421qu+Wz6IiKFLGWLeJVezxqO7ZGM9urvCv2ab5jvjcIP+feKm4FZdPaxECQKnIxHMam4BotNzekNonwl00Z2RM+jB7yv83R318Pes6KQKtHPL6R0XsueVsekEyDaatAChkxhbEahWUZifh6/dDCZMEfhAoA/477YF00SalA+UoOQGmmlAHT7451SAcCcL0g4OmpqHlsQfAZ4DXgRTYSA5ufeM3HsgB7A5qfcyKT74xB8v3k/lVCgkW3eQ/zfEyxroXbwH3uGfoGU5LcnnKyxyUzItxqgfXAUZdZ5fJ8HCTVBVFxTW/sOc/rFb8NOZj0azPv05n4/mXUJ/9jy32/PoT3/YG/WO6c9xfzia+fTnsTcY9md2ixMkR6OR16M5M2/m99hPZzQj3twGGComna+3BFJ8lbPRHP9HsV3ijcgx52w2G3Eejv35sctTx8fjIf3pDWZdf8B+9sfeuJ9zNh/Nhj5jx5/5s2PG/pj4LhnknO3l3HkkDN/CLApYGk1YhjZpgZnJPLgqzVl86HF+isgzwBc9HLU9UDKdtrQsPoUJUujZAr+YyzCCrbxDFwbBnTMKLVkhNxDahIFLO/4QeteNweBrDBo07cmewZCgKGiII8oN+4H/NSf1NEddjWapYXG66ATntNnHRW6nF92WVfw3bSkZDk11yhmfBWPaLM3jVKF3MEIBwwBU2csnnLbe5D3u9iveS4mUgFJGmyB95b5ik8em3NBLEpc7vvuqz0AOB4XeYEDlAf82JcoPUiKI2JFE7Ezl8bdSLblinJIq67Gqypo2NRe8mJqDDl+AHmAic6sImmkQVQVyMQQWYD7eskq6K4vUNEx/LnXVjsyfrjkmYuf+6nlYlrGses21H6S67u9CdYVETUqoaj717bG+qTpbGp2uIj5w+Tjx5WMY7/Imqr8j4M4xCqb5PQ9tp/dV930V/mCVP1jptQ1sW6YjtT/nXpjVmVPZ6yxab67OU6t7b3V2f3/q7NbJvftgdd6brEGdsuoKt0UouKnqNp/j39/3MBZYIqHk4R/VBAzWJDiVTEpR3ebB9qLNsrY5kpa9dGGykPHosF0Kv32SR6m2KpF9Qfvaqex2chDgZwVswNBK/C/e5+f3PvWBsCJuYpJRHrZp0bhNSwrc3F8FOxsnnXvP4kq7xEnsazfJ0qe3zxE0n5hTucgTLpTfaFpOAbX1WtaROQcEO9gKMeQQffz3CAxoK+SAQ1KM/kMwjvDfwadhwr9DtJdi9inHmIPoBufxINC5C1qSGyvzgr+2HOsrLR5ZeE85dpasyBYrjMiVmwU3JF9CPC1WOAXMwr0Cd3vlK0sQus8mxWtw1TfH6WB/TLviZiddhkHWsG3N1WNIYrHytGRYPMfkNCDo5RJhoUwd70IhPJX7eJTyMqE7JC4TsiRull7Gc1ytXoWhMo4bAn8K2Yl1cBDoYx5v4WkG3LSsNPAJ54FzyViWAoIlGQLgK7pNh0sHYJsqcC452nWXK6PzVO5WDME1JjKNLWs3UVWM13Xl8mDjfUqloFqxRv+A9oUl76AYQbhyilGkpC35d6WHoDlDoj01672W8ihacpWqBj+pX7/Qqndh9v2KSZ4xw+AQTiuqa1CwXH9jxyO7RXKwiwF/8TQ+p6dh3Cuysxh2F2/FUgq6EWwp7knkv4j8wCNpQ49kBywdf6RgJxRV24AkOvQL3iXQTlisJQGSBT2BCVtYFuRc2OjH2FPqycCoQxI3tGVjo2V0lqv0WiBQRm1j7KlMUvfYBev5phlezQuO6q6D1J5O1YFPAJ9ZHCr9ECwvaT9ka0s9ErvvHt0ZwDcnajKJfEh8Z56N84JP71suwzNR3YZjtU1l4c4kXK+q53sXfTDl6YtmpQWRMnsU8R41shsLN/1AfCteZU37QWxehnH8YbUscSvWb6xvvrG+4rjBVRQnOEGkvWWNdqr5KleHmWq6mqUZ+Gi07eYmyHi7pGvV9lSbLgpOVdSqmWOJw1X0IYo/Rgp7FV6DhCcXVjUu7SJ7bIIm0aN3Z262HUQxt9nz7U0np/mQFlC2q3upjRWpsb6r1rbqrL6FVOhrUzURkYs7/eO5rU0q4pB0SJLEScP+mfEi9/02H0eMO7v4LgBWwPs4iMTcQ9lg+gSU/HZJvNIq7SU4f+7tz1EWhH9lG74bPsHYH92e3LJcCiZkJ41fbHt4kN3+NEtJckP37vBtsCRJCUUTWY0GuLhJQNJ8H7NQF0+/6E47gYT4BksH++uWBx++E/ulm3zAXfdnlsRu5+8rkty+BXfWy+LkSRg2bH3ntix5VreGvCohZkIkxPaiFaZaEIB0ErKIb0ijabLkkoA6fpBCHSL0PHRVtqy7Tb5TGGqRZk8imDYgf98l7oJIW8DNtGP2Q1ae8GPM27pnqyD0n/BN5t8FV6tE07xHQyWiztvsROXvclfqKouyXd6TP+E5zWMWwcYoD37JW13plvqnbkqG/QJISizBPmdxIgWUpsmQdNB6CeOSTlnLkHHw+AFGDHy6B53DS4ky7Fq0XQm0SJMhbw2Qt0bINIQBwDeAaxkyjhpCK1A8bRuI5OIGa+gGX5ME94AUCEqyJskV+Y462Lj48ELyfhWpVgFN9opOChsx1Vex10pWKQyVc3bGQBomKQDrUfQTD5KVlLdQaISDaCtZdiShkqjAzK5xYoK96rdscFhF6Wq5jJMMfCCfDvvN8l51aneX6CmphbJDIppVNmWh5ZKLQ5zh8n3wmOCtkgT6ajUxJUs6g+nKUxgtjpRv1C1MtiMiD7d6UlPbYIYFsEGTl18E/nN+kNUD8Z1v66f8K4kbvZ5qjbM4c8NncZhq9Y7Dv+EpKVpPZ1pksOqAVEtzOE0ApXobN9GlAAA4BbAUN+NBTNyOCzkdcJ3AoaI/qX8lwAqOmMgKv0SuW/H7MZA8oCKGsnCbf6ScochrzqjleIrB5ZJ7E3/UJAeW+wMJrq6zkuhudxXd7X1Ed/sJorutFx2vXPF7i+iKqkuyQ8RmhSnyIcfDoNpb7IF/oq4tFnq3qZaP1lnvICQNQ5JUufALKqcInI+pdABoT+fac2Halsrrdg0/9lYLaHgdLyFuRr4NCX41bAZq52eo6GeHnkLEjd+FaR5YPTwbIbYeKuDXVLI5PNWHGV4S6jOKi4ukZJ0pvHKqfAkVcht2z6csGjoJvt2aHkz5E9+VrQSJi4A9OJ84vKEKVddIWgIRI+ebYve7HG+uW7cybs4zonCeVd+Vb/KHAkvwC+ksgLSzv7Q1nG/t17eFlxaX2YYYevzjySw1lZhnqrsSckR3XYNIM0sMmxV1xuRbXsgoqYHvZqf7jVEguN8ed27rwjEv3nMHku5+fylI1pxtKv6Y+G5hqa1ChK1CKC2r23EGzcmu1VF4wsQDPL13WBzK2xoJ54tleELIYDpBZFDhfeyMnTxayKd0tlWq2xnsoJEaDbexOlTDyCn72hKB2TMc7QDn/YV+fFKK6dIOGGaUWg+sHLbISRycyTMEuSN+XN9RT3QO+RSiGE8kdwQtkPVNXW1DxW0Bf1sPz7uxYui38Ex1SKB3TvDE0p20eMK9Agabp1aCKz+lCQLgU8f5IpfXtGx0Isp4pnjDvAOw/kgXb6wT66uviuwKeoZt7HL8QBtb7rG9fa+0aKfYqGKD5p/MmfJv3Mgjz+JVlMm291CnintEwrbAvzmQlJs7OJha+DkUTDhEKmylVSsOm2q+RY8g83F+JjlquFZcWoTVPlXZlKuhVlOTncp7eyvv0DQV7nRaKDOg0qgQWrMev2LVWfucgb/1oXrdeKN1t5rtnOXiMG5G0EorlSWPG9ZeecuG0g+sq832IZMoOv+pMVnsog7OijlNpb1WWuu63lpRnGvVVtd1tqp8rLfY6brOSittdG200XW1jaGQ0ELNUmrWIRvNc2dzUdzXtW6U6zqj3KssoWqwVv8xdOGdMIjI3/ikxJnUAKZZEn8gFSv5VZSfuUsETv++chOyFfo/Yupq2Qtc3bV/0yF2r3Zg45U1bThhq+j5ObE2umQtQ++3DSbPfSzboDO1/vhHdFNxY0INhtSvVqFUDammcdSpHEedL+Pol3G0Zhw9/3zj6N5ug6dTMXg6XwbP3/Xgef4ZBk/6txwMK65TIZkSp2hE5KP4rW5WkjJwQDcFOZp8HVuwUBVpKyjdM0qG+022R8bkyF1pjVaPCOaSYPsSXuOyFl52cFt1xoHOk6XriGjwgQZgxECGV3w4E8PxehnrtgJLFQhucsa7axCf/8TUc7khFzH+4hKlKUDuqd3tuTJESdFtCctwFY4cy6z3wsSpuXQVZlLM+7NFVsTiMCNSxP+hge4UTrE4b2pInNKb3DN8oyBt1CgKWjvm88umShcpBGr8ppA+9vPnbLLUbpcrXre4lMNQZq3t3TwP6sX+KlylKjy9+inHoV8ynl6/4vNrQW9SOusaJ752h5eMeMj5Lu3RLy4hQwIHhjUf0VAwn7cPhnReLBap3aXBrOX+97c3EWh5sgjh02gitw83kdudTGSrv6rbiIxQayTlGj7MSBTE34+RiEvW9PhoyzKGOVucnSJoOa1doCzuw8R4vnRJULEBC0OXkf8MJpx+5TKgH9zYTfUixSCiVKV1OX1FhRV8f/qUsoLGVvzEAl9Bmt7ema9G2tKtqPbECJqvRO4Ai/cCf0evBaazeXExcAV0lrhRijvNf0qCKxYAyOIl9AHzKvowcL1O4iVJstuGHSzcK9JOCFp/EF3h/S50zw0IybebOxBot8Xto+1f43iBBJwdEfGqsTa9BTkFz4Ri9pdrW4ibqUMV9TvPDb3GjZs0tHLRa350Jy8Tb5ZrcSRQppRrYjdSDLyClriblp4iQJmBe2ALY2H4OytJBlfkRC9JRuGYWLXLi8+60PkFvz9CWXmtH92VOv4Nvb2xMyAL7GehvqLCFfT+Ei8lcuudyJV2MYThj+6MhPLmDr6qRNNfa+ce+NiC6+8LaBcUhobU6NJ5iJ/K+jlNERaEN1+9xYaE0l9K8TQZil2T/JRetY1w3Y4DNZAuTDZioQ2ze4+krsimosE72xy6LNvodobYuZaV2ESLYxlqk2iKsyib337PEK2PdqS6vtc0SUJuEeK+ZqPIQtkcqxshn7VI7U/XF9cUI+R0u1+DvT26C9D+qPmZ8Xh/ItV2Gye661pEKA3MYRT0GevkxSSCZphYQdgneP06ykzqESyjoXMrz2e0pSFLzc9bGNtdJGdufvO9VP/NJqX0mkmcAXnsNtvjrk+ubCNtQ7/MLSrBbt9YTKKMHztbnMlS29bDLf/eFlx2rGtN2GCcqkx2sM491YMGU/3xwSZRoG9tXRLoDiZUgp4VvX81kLCCOhjo+8XakI36roZUTDc3ynZMrRKNWLFKzZTl1l9QfmgXYaAgDcYM5o4CnUj1aOExK0yitDfNqh2ANPDyNF6TtNqD/6Q5QlFAxwvdNP0xSLMOeCzg6UbzGEXJ32DIPaf7hIc4GLvisHabKJ8FsV3FbzlTmuFrKwkS67tUG98ksaUrnLWyOI3Gu/+KHt0VbWRz8U7bwIOvX5RYqyqUPpshtUqGzKOols1e0rC1XLEVpywlDsD2iFVnu2tV3I4GkGZkKa9tmEXBxMlQtknNnlq2JidmNw+UE0PO5UQfRbG1zGoxcYAqMYnsSjFxAF1MeXJ5hoX+ar/rXePVKikfIIoRok6+jGhZdh4GTiHnGfQ3+EoJfauC3GSKw16SBQB0oKpXJGNJRehDs61qwD11pbp62zK/X0KnpDRvKS5UEf/XVQ59z7d4awN2RAS6S2EbLZNQ9Hb1EORyp4tRnBAofMYul1Ms+h+R0qEb6ezc/1iw43B26eA63vMNXS2e4scbbToDeqDewRto5BUCtQY/LfkwrvCfl70L6zElIXecKvHcovJLw+XzMuhOlTuCHbuBAtnYEUjZrKXbjiGHNnK7Z8jizdt2o9typqiV4UCVTIJkT7IsCWYwmDdsquGWrFotEDeP8Zmtzxfg4xSNI7cGUj+pF1CiN/8fMb6c1daJ8OxCAOZ8+4DO5+mx5HncLMcx8jeaKhtY2QnbSSCCMJPIK+iosJ4i1Z7oPGgD+efjATujMg94GZYGdL9wXXn6oRHaOVpXpsQ9sPwFp88YhxY0q1zMHEANz0l3rK9f0kfdYODQejoMAOHGKRrrR33BOPQU3/MKoqtnYQDMvFHOBEtLPhFMi94IcT26E5T4VKWdk6axFxDUO+OKkIgBSD5/xU6Q3MUVYToJR13i4SAdOkfQJlA51wZ4OdZqniKhtOVj2lgL6udyGWqktgiTLeHk+ML6zgvRsRTt5q1t7Gp+vso9C1Q0aJCiadEsDL432t3OAP0vp9MjC+2GiE+poaXaCY/floylEAK3Fjw45V3Ll+wZ9qI8mHRXDZzou0oUO5+vQlF5QbvOkov57TY75kRzBNzOs1WeEjfqqal5dlI82yGgaJVbojD6pbzGkcVLExYkF0jwoeBQuZ4Up2dzLJpR4NFPBVM8KVlGZTkFLvuWkDeGx1nqlh0kO6Ch37aliETujuoWHErhLvYaJo3A5QXg0opEHz6V+FYFfanrVCUod50JXw3aRixfQZCp8SCORI6lFPQ22ql0HJiAGrovdMtQWuxrYVa7WsIwQzD/uyRevOWeqfkYobI7AK868V+v2eUZfD2cpjY+BpEff+z45AZmGLRUCqTczF8Bg09dKZt79GIc61AuCj+3Fac+/rPMAZ7471dptmAxQK2gSqrVz+14afp6/Rc67aMvlCQpO8/f0H1oRQ4yllZ3laCjVwexFfTHWiWa1tfSexpnpetNPm1J13z/5b1ouqsstivURNcs0Q7lGsqb0B4bNVnVqLYsOEMbowVKt1bJHpayhMb4++hm3rVkt3qL4ZtSiB+otwxzs1og+kvMbbxrJIQGlOkTpY/uKgxv4y+xC9LCIiGf1pfKt1jZ/M4YLQTgXdNJYytHl0YjU38gj1ZaxaW4hfjBCi4HHqpK3ZQmkdVBC7WSNZ1bfntc6LN1Zok8u8oJkgVofZ33FM2lXhKH4VPFJ7ujg66pBHzLkHLQsmbk2r0J8BVXG+9XcqPM3uhlyBNpWs6LKIv/GpCPjTsDOtAMY+8DpETEBSMqCG7UQ+9leS7iVUotA2Wqx9DE1Ue4zxQkJ284xdAWq9x/tqzi4xelTxOYhhPQ+XSUNUv+Ui8N/MYRka6P1GZxVYBVN0HnpyEvb7Q7EvPNR3nNdimUCrqaPbYWrXjpW6abdNMsPRBdsYZWKiPTvJftRdADH9uLQHWoC2LmayeVsqsFXHNlTnOygyHoopaCE+AaAcRQz0OPifhioSynzEyTJR/kBMzIP4gogYbN003o9EKxojlAS5f5OLe61jffKCKTgQ80YDb1kjhWJ4iatPjWB1u57FaD4bNjbea429KxWi1Rdc12DJM3I5tGDnYq+qCq6KrKKhLZVGjql1xTP4jJeY2qfslVlUPLuvrBMJvXeKPttl5XYr34E5X1y6coq9y93ENXv9xDV8XiuF15BGSn4StmHoE+eu04xOw0wGxlhMbpto2hb8RRgweOpAK/PJyWT8UgwMMHx8tiVGSFTvTrVC9vyte/1nNhvtkjD4CqI4/9DDOIf4LBcs0at1yGrFemdN2J6c9ug5WlPpdlPlKknqjaJpOKg1glNP1RAPMNITvZ6sdrQkKTrYr+0Q0zKLcc2PJJmLm/4BSD4TY7LKXgCrH53fLPydwFw9FDosVD1CWnvYkPLxT5jLb26E2tpy/LpAyI+9o1dvnUAEYNhPsR+A53nLLn99PJiAqhiX4QuszRudXfVrv2mdXX7U8p9bSCXbzNpnSdjcqu8nmgnmA0HfWrrMep1d5akYNtFTmvqkgQ3asi7e0VMSxmyiR2nwOXfzxs6aU8qbzTI7m0YXoU9xc57ochXOtQs+HHWp/X0CSkQCtxXBZslgr7T32553OWZprLFvoy3Fa7dNM0uCEn7PWWjbImZtpFUadAYwCjdGEtC1T9FJHn/M2ez3ddrXq7kHzRqGHzKwtecCCRYoDLw8UyKEucqNczla/Q2/0SvdI1euoleOVr80r5jIVM3y1WfTFe+V4k+poEe7SBD/541wRYCW6JOJFeWQLL4aOudjI2P1j0gF2klUd4HM9d2lVQxemdWrAFjTgyndvdzmhAFpWwkscaRHh7R1udOZfWFcxHTczA6BzPQ3oY2GbXP1fD1p82Mkis8hyFDom1ymfmdq8zsHeK724LX1fVeZVhgVRJy7V2ZGOv7mxYfshLaQS19iGDVBmHDFOxAXlPW3T+Cw4LpUvMMefbsDQX39Hiq2otyMrOEU/S696T9y6XwIyHXywMymMTgKlmvpzp0H1eMBKJyoqbBjZk8a6Sfu3m7RK0CK60B3U8F7NYx6mDU7a/eySSL9YpF63Yu2MCVKdL+oCw0e2BjQC/H5Ng9f23WQWP9fzbTUIM/KpVVG1ZynVZvsOfn/G9428FaD7RSTmJ+XGG6/RPTIkIvZnsbZrUC8uuA+oCvInj7FXsk0azcx2nGcw251HaEQEocSMj/JycHoLnHCyz89PDLCEk9WAIaCerqH1NEnJ+ihvcLbrD6mx/Hoc+PuVxGQHl/fNTKqfzU7qwZKHncLbvXRPvA9Rh34hzmcVXVyGiHlIklTy97ePSnc1gFrxfl/VNmE3eu+sOFYTFr4Vv9FtN61///J/dTtf6v/8H/7741z//N32k91///F/w79T6lSTxSR8qDJTP8394LeTiYJofXZI1pPjE11jBc9xgKP6lh4oH5dXmQybuzuR5QuiXubloyNrbF/vnb8QIzBTW6XQE30ZlUU1yVYB5hoFHTfUw9jKStVPAcRf75/l7Vszcqc2wLzSZSR6ymNNr/GtNasLAOsDD2ziOxKsXnES25EEB7A2Vpz4adkYWS3QnkA5JEhAYOG54IKRwMdmrBf/x9qdXHTrBbSDBDj+3rdNjdbebSlsGaoiSv/phqe2gY3qzgrFbPHiBbXlLIxF1EWpUqrRfo5f3aQz2c7ePexb3T6z95ywsll98A149BvNgHmu5dIG+2bF+QMfykKXjViv2WgENBO63rH3pjQKk+MT45+wsB6XPICAkb0eYoT05QOn8+dunT/7M8KQ7DiDrbt/twj8OTBX2i6OtkHJxt4+7+hEbICCXtjT6CVBW98TqYyrd0g+p0FD3wdrhV38zhd+3Ci1M0U4EiWT1liDKkahTcZ8QlgRjHZaDXcL+wl3nv3kwEr7ZLGg/v7o5B8lvcM5TxNvFyMXFeMh2zR/RR1Wdoz777Dv0cziCnCH+n34ew88RPg/LPsfHANzFd+zY27Nd/D4aw18Oe2rWwYQRPpY4YAk9SOh1jyBhzB517eP+fSyhd8Qguj3cwU//YoUMEX2Mb+t1uzTBwX3+oy5NxYTRAFnGko4pxhH+7B3l30hs3M/BscQhVnpIn409Amp9pOg40ykagrLDHUQ0YlZVFShm1iK0xl/G2xc41GTQVvcxkFELdqGC0Es3qKmD0RSGWLx1hdYmEoTxbStjup2N/4os6LlUOJa5mW7kzoR3HPATholzfWhi/+w4RmqD1j+s53T7y4n17PXPVpcT++Yqm5hLYQzgYEW9JujE2OT4xMLFnH0r8HEwTy6Lvq3f87o9ZzwbOU5/SPwZcWeONzvu9p3umAwdzt7OA3njotvB+9HF/8Gu6IspZ7xjapp4vy/TwNmAOLPZeOZ6fZ+Mjoczx/EH4zlx3KOeMzIy/f+734OxcPB82tAxHVPfZ0h9H/jGFgvej9NxpigR6gA5wy8e0BcP6PfjAX3/p/Grm5FzmyR9pxc9O/zl1bcfXv8pIdd/Gf588OL4cPCn/o+ve7dvb0aDZNF/ebh88nF1018uv/9x9P7PvQ/jRXYT3ITH7vejg+jPr/oHT5Y/vr65eqgHpf+/wqNy0KNqWcy16te6Vk7ZtXLu5Vp9mpv2eZyyNnRRvf5odDTo9o764+PB0SB30oxZ23w2p9sdjIfH/YHTH4/H6E4pPpyR5heX7vfq0rWsbYU4X/zGz+Q3Dgfd0WDeH828PriQo5HbHx4djTxvOO577nHP/W1dMIuCEf+BrpgXLzCYe1nnH9E8AA/dZYqZ6LaZdKLR+IPFgT6Li1fLULUyf2NXkDv6F8D3cfeoP+iNoReBj6N+d+SMnSP6Me53j3vjodWypI68eFr+AlJ7Q5hy9PvQgRU9OYMf9AfDo9ERJzvqO8O+ZcnITqfbO3aOuwgPIMPR+BimLOwDesvjftfhJQ/AIscD6VF7ZLvfP+oeO8NjhtwdHOPsh8L3oRLHI4tR6ncHMA06PppCV8oCx8pUpkL6hob1yU4uu0XhhU+vcCv8sqF7NIJpT2846s37w9Fw3J2PemQwHBCUiEe7u92nejLw1imWDLy1M5hOuC8eewm430ZvXDjMKCsOxME1F3m/k4vzErP3S6jPwpguX4rvjocJtBj6GsOkyEkIdHMe+VuQXTcU9JzoPHGv+NkobU3oOf/8jkNgHQS0sq4gCGO+4VhlrlztWCV751Gce8opwydfjXp6+wKIC2y88HzCXgiNV4lHntMLBSpk+AfszfetA0tD59d1yGLJqXXmQZKKsmnNAKHILWYnG1UTRiFvm4z8P+G5qgk=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_9572198988a54a0b9731fab9e0c97e51\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>\n",
+    "\n",
+    "\n",
     "## Multi-process checkpointing\n",
     "\n",
-    "In multi-host/multi-process environment, you would want to restore your checkpoint as sharded across multiple devices. Check out [this section of the Flax scale-up guide](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-sharded-model-from-a-checkpoint) to learn how to derive a sharding tree and use it to load your checkpoint."
+    "In a multi-host/multi-process environment, you would want to restore your checkpoint as sharded across multiple devices. Check out the [Load sharded model from a checkpoint](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-sharded-model-from-a-checkpoint) section in the Flax [Scale up on multiple devices](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) guide to learn how to derive a sharding pytree and use it to load your checkpoint.\n",
+    "\n",
+    "> **Note:** JAX provides several ways to scale up your code on multiple hosts at the same time. This usually happens when the number of devices (CPU/GPU/TPU) is so large that different devices are managed by different hosts (CPU). Check out JAX’s [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html), [Using JAX in multi-host and multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html), [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html), and [Manual parallelism with `shard_map`](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html)."
    ]
   },
   {
@@ -351,19 +414,22 @@
    "source": [
     "## Other checkpointing features\n",
     "\n",
-    "This guide only use the simplest Orbax save/load API ([`orbax.checkpoint.StandardCheckpointer`](https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html#standardcheckpointer)) to showcase all the tricks on the Flax modeling side. Feel free to use other tools or libraries as you see fit.\n",
+    "This guide only uses the simplest [`orbax.checkpoint.StandardCheckpointer`](https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html#standardcheckpointer) API to show how to save and load on the Flax modeling side. Feel free to use other tools or libraries as you see fit.\n",
     "\n",
-    "Check out [the Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for other commonly used features, such as:\n",
+    "In addition, check out [the Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for other commonly used features, such as:\n",
     "\n",
-    "* [CheckpointManager](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) to track checkpoints from different steps\n",
+    "* [`CheckpointManager`](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) to track checkpoints from different steps.\n",
     "\n",
-    "* [Asynchronous checkpointing](https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html)\n",
+    "* [Asynchronous checkpointing](https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html).\n",
     "\n",
-    "* [Transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html): a way to modify pytree structure during loading time (not after loading time, as in this guide)"
+    "* [Orbax transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html): A way to modify pytree structure during loading time, instead of after loading time, which is demonstrated in this guide."
    ]
   }
  ],
  "metadata": {
+  "jupytext": {
+   "formats": "ipynb,md:myst"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
diff --git a/docs_nnx/guides/checkpointing.md b/docs_nnx/guides/checkpointing.md
index 1ae5978b2..fa98e6db8 100644
--- a/docs_nnx/guides/checkpointing.md
+++ b/docs_nnx/guides/checkpointing.md
@@ -1,17 +1,39 @@
+---
+jupytext:
+  formats: ipynb,md:myst
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.13.8
+---
+
 # Save and load checkpoints
 
-Flax does not actively maintain a library for saving and loading model checkpoints to disk. We recommend you to use external libraries like [Orbax](https://orbax.readthedocs.io/en/latest/index.html) to do it.
+This guide demonstrates how to save and load Flax NNX model checkpoints with [Orbax](https://orbax.readthedocs.io/).
 
-This guide will cover a few user-side scenarios of saving and loading Flax NNX model checkpoints using Orbax. The Orbax API we use here are for demonstration purposes only - please check out the [Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for the most up-to-date recommended API.
+> **Note:** The Flax team does not actively maintain a library for saving and loading model checkpoints to disk. Therefore, it is recommended you use external libraries like [Orbax](https://orbax.readthedocs.io/en/latest/index.html) to do it.
 
-> Note: Flax's legacy `flax.training.checkpoints` package is deprecated. Its documentation resides in the [legacy documentation site](https://flax-linen.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html).
+In this guide you will learn how to:
 
-### Setup
+* Save checkpoints.
+* Restore checkpoints.
+* Restore checkpoints if checkpoint structures differ. 
+* Perform multi-process checkpointing. 
+
+The Orbax API examples used throughout the guide are for demonstration purposes, and for the most up-to-date recommended APIs refer to the [Orbax website](https://orbax.readthedocs.io/).
 
-We first set up a checkpoint directory and an example NNX model.
+> **Note:** The Flax team recommends using [Orbax](https://orbax.readthedocs.io/en/latest/index.html) for saving and loading checkpoints to disk, as we do not actively maintain a library for these functionalities.
 
+> **Note:** If you are looking for Flax Linen's legacy `flax.training.checkpoints` package, it was deprecated in 2023 in favor of Orbax. The documentation resides on the [Flax Linen site](https://flax-linen.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html).
 
-```python
++++
+
+### Setup
+
+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 orbax.checkpoint as ocp
 import jax
@@ -21,8 +43,7 @@ import numpy as np
 ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')
 ```
 
-
-```python
+```{code-cell} ipython3
 class TwoLayerMLP(nnx.Module):
   def __init__(self, dim, rngs: nnx.Rngs):
     self.linear1 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)
@@ -32,7 +53,7 @@ class TwoLayerMLP(nnx.Module):
     x = self.linear1(x)
     return self.linear2(x)
 
-# Create this model and show we can run it
+# Instantiate the model and show we can run it.
 model = TwoLayerMLP(4, rngs=nnx.Rngs(0))
 x = jax.random.normal(jax.random.key(42), (3, 4))
 assert model(x).shape == (3, 4)
@@ -40,12 +61,11 @@ assert model(x).shape == (3, 4)
 
 ## Save checkpoints
 
-JAX checkpointing libraries save [Pytrees](https://jax.readthedocs.io/en/latest/pytrees.html), which is a pure, possibly nested container of JAX arrays (or, "tensors" as some other frameworks would put it). In the context of machine learning, the checkpoint is usually a pytree of model parameters and other data such as optimizer states.
-
-In Flax NNX, you can obtain such a pytree from an [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) by calling [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split) and pick up the returned [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State).
+JAX checkpointing libraries, such as Orbax, can save and load any given JAX [pytree](https://jax.readthedocs.io/en/latest/pytrees.html), which is a pure, possibly nested container of [`jax.Array`s)](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) (or, "tensors" as some other frameworks would put it). In the context of machine learning, the checkpoint is usually a pytree of model parameters and other data, such as optimizer states.
 
+In Flax NNX, you can obtain such a pytree from an [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) by calling [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), and picking up the returned [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State).
 
-```python
+```{code-cell} ipython3
 _, state = nnx.split(model)
 nnx.display(state)
 
@@ -53,7 +73,6 @@ checkpointer = ocp.StandardCheckpointer()
 checkpointer.save(ckpt_dir / 'state', state)
 ```
 
-
 <script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: "open"}); this.defns = {}; this.state = {}; } } customElements.define("treescope-container", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = ""; fn.call(this); this.remove(); }; const child = this.querySelector("script"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector("script")); }).observe(this, {childList: true}); } } } customElements.define("treescope-run-here", RunHere); } })(); </script> <treescope-container class="treescope_out_5ecb69e45849446495cac4e3367e5c4f" ></treescope-container> <treescope-run-here><script type="application/octet-stream"> const root = ( Array.from(document.getElementsByClassName( "treescope_out_5ecb69e45849446495cac4e3367e5c4f")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement("span"); msg.style = "color: #cccccc; font-family: monospace;"; msg.textContent = "(Loading...)"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: "1000px"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll("script")) { let newScript = document.createElement("script"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream("deflate") ).pipeThrough( new TextDecoderStream("utf-8") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(""); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style="display:none"> <script type="application/octet-stream" >eNrtWgtT27oS/itqOnNILsRJnBcJj7lOyIsWKIQWyjlncmVbdkQc29hKQjjDf78rOe+YFFoK5TQw0xBppV3tS9+uuuuzoUX2JeYR4muOS1qe4zD0D3IdnzLq2EXkEQsz2ic7yHBsFjdwl1rDIuo6tuO7WIPxQZsyEhdfisj1YMSiPouLreNs6MKo7dgwrGKtY3pOz9bjmmM5XjFYuoNG31QLCGA/qrN2ERmUAZnNiM12kIt1ndpm3CIGKyJZa3MmNom3CTXbMJKSsnwbm2EKMk+Wjf6I96lPVWpRBpLjHnMmtHFqM4/aPtXiPr0jwexI3PvdRKCe3Yl64l7PBp4ejPmaR12G+Pn2NrDrWlTDXGMJR2OEn94juLuxH43G9vZBocDPZ0gnhu2jPcTa1JdMws5A28eOTqIxqe34TBLzcDTCUMslNj+yovFd+aI//w6bqWNbtwhM2z3L2gk4SCBm03FsGI0OHK8TQ7MyOBcwxKfmhhnV+KBLPMPxutjWiGQ7g2hM2BcYRJdmUDxYtIvScgz2oQaKLkgtWcQ2WRvt7aEkJ1kpukdYz7NB74hYPpkK1u7ZXLLFrf02NRiXTxDwP+7h9wEOUfAqW3cGkkduesRnik27wlxVD3dJNNBJjO+xs8TI7fntQI07IWccs9gLjrHilI+XgUsRGJI5pmkFUdkSkQPe6vK9+Aix2BYifXDwkSW5dOK71CFDrvSIF+ECjYglzcK+/xGCc7RvNDLZs9UFN4yMmd/HQJ/g/sLH93cTYQGg0z4SG+5F5tNHBDGswknJ7V4kGYHQ9dgyiWODiKAMG6ZWBUO4BqJ8zfjsEQjGII2JPNLCquqRvvAfkVbe57ZlnEzCqUYEmtPtwsIZCix++OEXSHDRdli02Hb6xIuF0M+TtwZtYrfIrQsmJ7pYKhmOpWMVTmDD0Ypt7Ef3LawSa39+phWcM2CntYnWIXoshv4TQw9zBQ2CW+ozFEn4MYwphc+Tmzm7RTqVTWU5AWiXeB7RWy5kXdIGWYg3S5jnv6OkL9RbRJRhyHJ88ZzoD1wXoMuwEwK1Tn1gOhxfC4uEaB8JDRWLKoGMQ2ak0sTPTii/4GqIp/jdMLpCwOYTXtQWF4ZqOfyOeZCnsPQyZx17HZ9gE7zYXl79TJaeyMCXhi8a089JKO64Itr4S86q2sZrije/6EEhcy8gJLcjZ9zzfG5A14GbnnghfKn/fGxFKAhGcZGb/Id8/Hm4To/HyC1b5iJRv2VQz2ctx25x9w8JrVWhJMlZHk2hpkI/LH5g8UUR+am62DMBmwViiIC+/0FukA/dodpjDEBRWAKaToc5bQRFFqhAkQB4w4n/IqmMHlnAyxtHGLyCYgs1h13VsXx00mP8vDoqByvh0x1CYMQHRO0A9g0ybxfuvTbkcECmNoPlFPtEnyDm9yTJf3eW3TxYLSBtUiqQ7uIpg/gIOUV4upuulAbYb2lwD4BiJ+uxweZuj3GeXsVzYc08y1nVoz72ovG4jhmOYxsMK0BTbHaYM+FI0MP22JvFtijlIwIaA5gfd3rsaUeZSACGoUR/Ny+JYIne0a7reAzbS3urntMBMMBHpsno29qdWTajz7GZ7yWOq0AwvaUBKtc9Yo9Ena97YM95wmfDIpPQGV2kc6GqYUuLQnEGlUHKvRVgUvIZ5usn8v40SVSAQpB3A0l0h8HZuRSzyrvpYcsGrN2CKtWgt7DJXJhsizCBwgJzXDTAng2B1xon9rEtDANrqXQIoQvI/J9JpeqNClOevUZKGg3Fk5JIq9NyuCiKWezFTQ/rFMwWTaWzOjG3kAMebRKUBOlyWnsr8HBA1DxfiCE00vKSKEuJ9dkUP2fy8WHuJYF9QUkWdiFFfRvmPT2HP8whAHeCxxwEf4DmOeQIYzE66DxBC3D4tHQLU8vD5M9jsy200OCR5ss+9EiBF7U8MUUL1oEI7IeXveh5Hyn+nFFXG/IxXB+losfo8jHMvnWo5+hw8aYDeqd4Hh5Khud0o7qj9Xg1LPHk6kt9bPWIH43FJN/pkqhIubxdwT+lAB7xVsUjAVJkA8VQbNIc8tuEMN5BIgNUbjab/DRNPsb7QWJS8oiocJtDW4v+778jUKaRcfJ/OkCbLYlt3hGzRmODUTMyw5sNvqcVUc+zohwtFPl8YuAYhryjAh7JZbb0ZKF2ZColRfw0ThXFEX+Vzgbwb72qKBVl1U+pqyhmx/mgNyql8uCropx/LR8qR41SWamat436xzbzS0eUmOnqwaX8sZH72m+6PfrpKHueOrxsnH056l8c3bFPw2q1vHlhds5p6SDZpgenvcOKXrtO1tWE0W/o7s2HXPvmgtLT3pFda9eNz0z5nCsdexml2rA7lZz2udezN8+yN5rfGfSNqpW4uTUrzrapHg5q26m6krCVs+xHzztMnW2ad8kzPakcGinzOF8e1K5lM+kMe2f5fLeSyg3ql4UT03TJeWeYIQ31Lqup3kmNYcU8bRwPDrA/9E97jcblRaU6UD6duo2v+udEYtPMn+cv0yxpfPh0o/SzsOdH5TivHA2Urnl31tzsXTVJ5fJWNnLa3XHmrD7M9krKh7vStVt107R+Wq4kr3qfMs28bZQ+VurVo65CN7f7Fbltp9r5TfXL4PJ6UPf6B7XPZfvaqFRMtnmiXVlWPlsoHw5K2+1C5uio1kzXrhSz28hel04L7LxG6oVKqdSopQ/MzFniqzZUlRrY9MuHhHJawwo5KltK/a5yYl4xM1f6ZJ6cNA5KHXqaJdXSZblU1WjSbXuOa4NvuFeVg9RdqtM0ygZrDz/YdR1X/bqRPO7WKse5kq7cfPniYuY3r7q6jmlBNu4Kmc/0+ibndr3cifO13KRerds/rKWbF810tSJrpVPjfLNuOW4tU/UHWWze5LbpFWkeW+6FXao3iH7kkd7FTa3cTV1UvU6zeZuVcxcX/kABiWJIdINZdEO49QbHIf+DfybRj3XHBQw2DUnRw5YkaQXFVhCzf8Neq7uCbdFUFTA5QPCwN7iHraFoAKTnW94QgucOD18gGwFtPuZDeuBb8MqBw208wJQhG/epiZnjSbCzqzrY06WBRxk5hyI7Ot0LDjvaa9pXBfAXjcyUFbyjClzOaZdA/REdt9yX1nmkCzXC0tL7LSQnk0mB8SD5AtyLigI5nO9M7RCZCsdbA+MMxpvQEfQeVTG1ILExB3HidyKzAa6zAR9DNqagM4J1Xvpszupu1B3+Rl+Yl1HjxvB8j20RkUb2d4N7fZfabm9000TETa46t5HQTUaXPkwGFz4IIRbP852/aSP7f1iMiw0Uq+nmJheqhci+YeFbybZveU3DiEVVabxn8NHkw9F/dhOjU81utlS3RVbNL0zONnPhMO9v5fxOUDGkgi8jAYroYe0/q7JDz7gxLWM3kGOXuRvvbTwxjkVnP7aBJjX2XuTPjdFZN/6OIHHv7kVmivAi+uOm57CdGbJgYActNQ3gmheeDJmyDX+P1Lw/b8afazeVYv9FjPYqkdHHAKFAEn8hNL6MxoMQeSXv+XODK//RXjQmf0VvEvZ+SfdY8ZwVCXGQ6dx79KAXhdaVP+xRn7CHu5OhkIy7SvhniOrglTEIIDGBgqgOEVyIGsz+Ya4It3BhXzAwJO5AT4wOseYZQ2TOAyP7CNDKAbgEr8MwE30ylMtuI8dAic8+8fwE7Q/vQHtmIjlIJrj2E2oPjJZotYhOmfD7lsTH40mpICXj7jAd5+UvSDkM6MFaiam13OGCdCs+tpaNOOkqgOxrS69W4qqQE22DvRDXGL16R/aPwYiP0PDWL2wFccinmkEselY7BB+x1UBhJnO9ced/2fv/8S7/TezYIZ5NrDV6fCX0GKj/Cf4zXrBGkGsE+W9HkGNffxqymFu1RpFv5iJ9fWv/EJIMv5+ovheZ/LfIbaKmUljN5PF2MpPZTmO9kMxkk9ntgpohqZyxENJh/50ysn8NkSxexpBhOZil5WhmC2Vi3z7D8utdqJCGpqeS6XQ+l8mkM9taDhs52cjn5TTJpnAmk4vs/+zUH5Y3w3XB09wD+ljGlt9SzffD+5XY5GVj52nof37Zz82VX6gP1ym9EzuiNtUhohC10fQtmT8yf5f//maKfsZ6680Vua+Fl8M+7t/qXfsyTx7PWKQGIsm/0cOU/DgryeuHqXVrIdR7ntKYktcPU+u2wm/RVpC/47lCXj9MvUGY8wtYev0wNafRR1dM8vph6q3e/+uHqX8LenxKoS2vH6bWCPJ3QpBPfqqQ1w9Tb/cifX1r/+yHKVnOYbWQklM5OZ3JGakCTuVVdTtfkFN6IZnXf4mHqWxaxnpSy+t5PZXJZ7NqKl/Q8oR/z+XknL5+mPq1Hqbk73svkdcPU29V0b/pw9Sr4uVVD1NvTIkv8+ARqrDndM8JpU77+/8HnO2ZvA==</script> <treescope-run-here><script type="application/octet-stream"> const root = ( Array.from(document.getElementsByClassName( "treescope_out_5ecb69e45849446495cac4e3367e5c4f")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type="application/octet-stream"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>
 
 
@@ -63,15 +82,14 @@ checkpointer.save(ckpt_dir / 'state', state)
 
 ## Restore checkpoints
 
-Note that you saved the checkpoint as a class of `nnx.State`, which is also nested with `nnx.VariableState` and `nnx.Param` Python classes. At restoration time, you need to have these classes ready in your runtime, and tell the checkpointing library to restore your pytree back to that structure.
-
-You can achive this by creating an abstract NNX model (without allocating any memory for arrays) and show its abstract variable state to the checkpointing library.
+Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.VariableState`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.VariableState) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.
 
-Once you have the state, use `nnx.merge` to obtain your NNX model and use it as usual.
+At checkpoint restoration time, you need to have these classes ready in your runtime, and instruct the checkpointing library (Orbax) to restore your pytree back to that structure. This can be achieved as follows:
+- First, create an abstract Flax NNX model (without allocating any memory for arrays), and show its abstract variable state to the checkpointing library.
+- Once you have the state, use [`nnx.merge`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge) to obtain your Flax NNX model, and use it as usual.
 
-
-```python
-# Restore the checkpoint back to its `nnx.State` structure - need an abstract reference
+```{code-cell} ipython3
+# Restore the checkpoint back to its `nnx.State` structure - need an abstract reference.
 abstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))
 graphdef, abstract_state = nnx.split(abstract_model)
 print('The abstract NNX state (all leaves are abstract arrays):')
@@ -82,7 +100,7 @@ jax.tree.map(np.testing.assert_array_equal, state, state_restored)
 print('NNX State restored: ')
 nnx.display(state_restored)
 
-# The model is good to use as ever!
+# The model is now good to use!
 model = nnx.merge(graphdef, state_restored)
 assert model(x).shape == (3, 4)
 ```
@@ -111,25 +129,23 @@ assert model(x).shape == (3, 4)
 
 ## Save and restore as pure dictionaries
 
-You might prefer to work with Python built-in container types when interacting with checkpoint libraries. In that case, you can use the `nnx.State.to_pure_dict` and `nnx.State.replace_by_pure_dict` API to convert an `nnx.State` to and from pure nested dictionaries.
-
+When interacting with checkpoint libraries (like Orbax), you may prefer to work with Python built-in container types. In this case, you can use the [`nnx.State.to_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L170) and [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179) API to convert an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) to and from pure nested dictionaries.
 
-```python
+```{code-cell} ipython3
 # Save as pure dict
 pure_dict_state = state.to_pure_dict()
 nnx.display(pure_dict_state)
 checkpointer.save(ckpt_dir / 'pure_dict', pure_dict_state)
 
-# Restore as pure dict
+# Restore as a pure dictionary.
 restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')
 abstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))
 graphdef, abstract_state = nnx.split(abstract_model)
 abstract_state.replace_by_pure_dict(restored_pure_dict)
 model = nnx.merge(graphdef, abstract_state)
-assert model(x).shape == (3, 4)  # the model still works!
+assert model(x).shape == (3, 4)  # The model still works!
 ```
 
-
 <script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: "open"}); this.defns = {}; this.state = {}; } } customElements.define("treescope-container", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = ""; fn.call(this); this.remove(); }; const child = this.querySelector("script"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector("script")); }).observe(this, {childList: true}); } } } customElements.define("treescope-run-here", RunHere); } })(); </script> <treescope-container class="treescope_out_277cbde407034bf9addfa8e4969bce7e" ></treescope-container> <treescope-run-here><script type="application/octet-stream"> const root = ( Array.from(document.getElementsByClassName( "treescope_out_277cbde407034bf9addfa8e4969bce7e")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement("span"); msg.style = "color: #cccccc; font-family: monospace;"; msg.textContent = "(Loading...)"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: "1000px"}); window.setTimeout(() => { observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll("script")) { let newScript = document.createElement("script"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream("deflate") ).pipeThrough( new TextDecoderStream("utf-8") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(""); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style="display:none"> <script type="application/octet-stream" >eNrtWQtT4sgW/is9maoBrhITnoJK3YC8nFFHcUbH3S23STpJS+jETgPilv/9nk6CPETH2XHn7guqBDqnz7PPOd9pd0Mx9UhNFZyQ0PQDcsV9X6DfUOCHVFCfVREnHhZ0THaQ7TORtfGQetMqGvrMDwNswvrEpYJkox9VFHBY8WgoshHrrJgGsMp8Bst9bA4c7o+YlTV9z+fVeOsOSn71PSAAftQSbhXZVAAZE4SJHRRgy6LMyXrEFlWUM10phJGsS6jjwoquFiUbJjAFnR+2JV+yYxrSPvWoAM3xSPgPtFnKBKcspGY2pHckfpqoe7+7Fbtn98E9WT5iIJPDWmhyGggk7dtL4SDwqImlx7Z8UxBpPSd4mKql05m9GjgU5IUCWcRmIdpDwqWh6hBxCt4+8i2SzqiuHwo1eg6mEYGuAsKkyYYpucpNP/2y7kkHM8sj8JiNPG8nlqCCmj3fZ7Canvh8kEGLOvjnsCQfLS0LasrFgHDb50PMTKIyf5LORPEFAelHT1A23rSL8rkM8KE2Sq9orXqEOcJFe3tIkyTPqs6JGHEGfkfEC8lcMXfEpGarrEOX2kLqFxHIL/fwfkJCGk4Vs/yJysnNiITCYHQYhavF8ZCkY59kJI+dR4KCUejGbtxZY+NMxF5sxjNWvlwHqUUcSOE7jhdn5VWUOXBaA8lLrhBPbCIyhgOeRFJqF/1WB2Qqna5wRSqUEKumh8PwAyRnwjetPPC8GsIxVGbC7zPgTzj+0Rmv7W6tSwCLjlHEcE9ZLh8KErgPlpLbPUVTIHW5eEziM1ARnMHg0XPJsN4DablnZrsCyRiXsaiOXJn+cAjr0QGK6spbHL2kbSskuMp8ka66/pjwzBr6hBxUhfhbCxw1eNn2nCKUVcRZFJnXi3pREoAZhHNiXQVQ3ojrexbhi4Rl+U6qa2RHFVGBoZzIzTaQ4z7Yz8AxT9RlsGqJ7Cp2GVBbNASh01n9XSVENeThPvGq1T6B1CYLWpnRa2etvLgGZ3VZhJNare3MZVEWVea+58ti/qTMyOePJVuYD0KCHTgu7PHuKFwrSy4O07WIZ22tH+IQmy4xB8TKZNB/MnMd5Nb1m2b0SxpGzaSKUj/nin0z9f9Ub3nTk0qWfoCSMo5S8IiHMoCBDy2V8DVyafh6YqNUiARloyIQPnXGX0fq3DxBbsVjKSoNr2zKQ3Hlsyt5/Nek1nOppOaKMpvWhgp9t/pxxFdVlFYNMXcABMVqRAl9/53SoB4G0/5ICEAf6wrQ/PG6Q6sgZYUKHAnIcj3xz0QvWMoKME0dYjgVFHuoNx32fS9ExyMh7bVQI94Jn8EUEiM7If0BgMy48g6hwbhQwwECMgHbKQ6J9QBN3xJNvnceH/N4d4QdNbVChqtWxvmxxor15W6+U53g8MqEPgCOfdiPbbHUPWZ1+jmZK3uWRS66Ho0xT2ezFhY4ixkENkInmcVlKURCLo7Z7DRHbJEeIgIeAzyd9Ufi20x50AACQ4n1ZlmTSCR6Q4eBzwVmj3j3uT8g7EquzIvR1727sG3Bn7Mw36sSwIBi1pUJ8NfihCWqLg8YwHOZ8JWq60LqJI10KVVN7JlpmIIAguvBbYTa1FBguf9B3z9Mkz5AIai7sSaWL8B2qcWi80IXS8QzwZxBSl3NSvbMy7aNTT2/hjAAcPvbw7DHk9lO1qXE/GQpq6lRwZxPlNVoHsQ863BsUQhIWs8XLeJsIh/OqkOQBulZMt3N+OwCKJWVIFpCif8eqfKoZL6aS5eCOTPmXp240Yn0PBxA8fk6gPv26vy0hBi2RTIiInILWWo9SfMaeqwTkRgaE8wnnnWuWB4p1OVRBj3DYdVUtjRcrSd8ibAn+bzifYIc8dAbg3M8VW3uD2FsNEdyllFlIwrVMfZgtkxnMmrow1AZtSc5HMpPNe6RcjB8YZdUUiiDMg+jeOgSIuS8Tiao0ev1pDU9uSan7+ghzLbRmNObMjP963+TzmzKjvc7u/TiXMTk/YOXrE2Sq5+CBiUn5GYVjbiXli2jKp9vTXzbzu30oSmVCpuWVmkfOkbdiF7dE8Pwo2/10wn87bQMo2k896oPDcMZ+O+tbrPemHwxjLMvjQPjsFtvGC3nttv54IqwfkiJk2/tX+Q+dEtfxr1gRD8eFs/0g4vu6efD8fnhnfg4bbUaG+fO4IzW9zWX7p+MDppW+1rr9LfscdcKbt6X3JtzSk9Gh6ztduxPwvhUqh/xgtHqskGzZH4ajdjGafHGDAeTsd3ytm5unaa/7fQPJu1tvWNsMeO0+IHzA/10w7nTTi3NOLB156jcmLSvc47mT0en5fKwqZcmnYvKseME5GwwLZBu/65o9vlxW2DDOekeTfZxOA1PRt3uxXmzNTE+ngTdL9anra0Np3xWvsgLzX7/8cYYF4HnB+OobBxOjKFzd9rbGF32SPPiNmeXzLujwmlnWhzVjfd39eugFeRp56TR1C5HHwu9MrPrH5qd1uHQoBvb42bOZbpb3uh/nlxcTzp8vN/+1GDXdrPpiI1j89LzysVK42BS33YrhcPDdi/fvjScYbd4XT+piLM26VSa9Xq3nd93CqdbX8xp32hDTD+/3zJO2tgghw3P6Nw1j51L4ZTqH53j4+5+fUBPiqRVv2jUWybVApf7AYOzEVw29/U7fdCzG7Zwp+9Zx8KtsGNrR8N286hUt4ybz58DLMLe5dCyMK3k7LtK4RO9vikFQ1469r80epS3h+ODdr533su3mjmzfmKfbXQ8P2gXWuGkiJ2b0ja9JL0jLzhn9U6XWIecjM5v2o2hft7ig17vtpgrnZ+HEwM0yqDo7k2kU9GxTsmW9Sv8ech+bPkBNOJ5SkY3hqqqPkOxGefsL8Dr+TsYN7rCirBSDOOANxwPZqJ0jKaWLxghBc98mb5AlqAtuRZCeZAsJHyUmAtPMBWI4TF1sPC5CpyDvo+5pU44FeQMJq30nBcYm/Ca32IBTkgrC9hS3l+BlDM6JABC07MLzkf7OBkCUHy09X4T5TRNi+AAFF9ABuloSlovdwFAKnPl5Hw4q2Dyyk9Bb1ELUw8Km/CRJH4TVTaAAGyEPajGFHxGsCXx78ai75K7uK/cwkksPbuGW75oWQUvSm037tC7lAWjpNMoUU/u+7fKWiZJ+4aHcesGJaLNtd92txJmiwo8wszKc89XHi5epCm1d29vc+WdGNPp8Q/wBJDXquhpo7/HxudtS81HhxTyWUOemr3UN6ZNdG2ZSaGHuWZP+SmV2Jj6RUFRm9tTFgafKnp3M/LFzgJZvADj6+qgBl01OjhQmFz4nri3lrht9vGHxqsPU/NzwVq6TlVqR4CLEsLNJbplmKrU0DrtlzCk8qOj9VNKGvviqM3IvzN6L43lC2I1IJwR77loLSPoGUtq7SkPd9rYrJRxv1wh+VylULTzFd3WyxVNK1cqhGwXV/VYdxeu1K7xrRohWmR7Phb5XLqwiQqZr9v8GHWvVbKSL1swDfY1oucKpl3CBV3D2xbRSnYu39fzSu0Prihri+V6X7zzxM4T/njniJ0XH4eV5Nj8xgR6IvY/MLvi4/kN+TXb8Ir1cemfRVCFoJd/piG0bXoXcUQutSDroIWj+dwnB8LfdWb/Ac595uP+r9oIfkzbfsXCH6uU+weAqtzLopP7F1T9GXIp922gKve3BVXE1i29tJ2HQV8rFPM6IKuyTQr2dsXMl4p9/U8BqvSSXsnlynqltL1dsLXt7XLeKtsAsHLlMraK2r+g6s8FqnLf2vdz/4Kqv4ZzXwCqXvMU/22a9u912MMOi45r/wOOvCUk</script> <treescope-run-here><script type="application/octet-stream"> const root = ( Array.from(document.getElementsByClassName( "treescope_out_277cbde407034bf9addfa8e4969bce7e")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type="application/octet-stream"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>
 
 
@@ -142,12 +158,11 @@ assert model(x).shape == (3, 4)  # the model still works!
 
 ## Restore when checkpoint structures differ
 
-The ability to load a checkpoint as a pure nested dictionary can come in handy when you want to load some outdated checkpoints that no longer matches with your current model code. Check out this simple example below.
-
-This pattern also works if you saved the checkpoint as `nnx.State` instead of pure dictionary. Check out the [checkpoint surgery guide](https://flax.readthedocs.io/en/latest/guides/surgery.html#checkpoint-surgery) for a code example. The only difference is you need to re-process your raw diciontary a bit before calling `restore_from_pure_dict`.
+The ability to load a checkpoint as a pure nested dictionary can come in handy when you want to load some outdated checkpoints that no longer match with your current model code. Check out this simple example below.
 
+This pattern also works if you save the checkpoint as an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) instead of a pure dictionary. Check out the [Checkpoint surgery section](https://flax.readthedocs.io/en/latest/guides/surgery.html#checkpoint-surgery) of the [Model Surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) guide for an example with code. The only difference is you need to reprocess your raw dictionary a bit before calling [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179).
 
-```python
+```{code-cell} ipython3
 class ModifiedTwoLayerMLP(nnx.Module):
   """A modified version of TwoLayerMLP, which requires bias arrays."""
   def __init__(self, dim, rngs: nnx.Rngs):
@@ -158,17 +173,17 @@ class ModifiedTwoLayerMLP(nnx.Module):
     x = self.linear1(x)
     return self.linear2(x)
 
-# Accomodate your old checkpoint to the new code...
+# Accommodate your old checkpoint to the new code.
 restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')
 restored_pure_dict['linear1']['bias'] = jnp.zeros((4,))
 restored_pure_dict['linear2']['bias'] = jnp.zeros((4,))
 
-# Same restore code as above
+# Same restore code as above.
 abstract_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))
 graphdef, abstract_state = nnx.split(abstract_model)
 abstract_state.replace_by_pure_dict(restored_pure_dict)
 model = nnx.merge(graphdef, abstract_state)
-assert model(x).shape == (3, 4)  # the new model works!
+assert model(x).shape == (3, 4)  # The new model works!
 
 nnx.display(model.linear1)
 ```
@@ -186,16 +201,20 @@ nnx.display(model.linear1)
 
 ## Multi-process checkpointing
 
-In multi-host/multi-process environment, you would want to restore your checkpoint as sharded across multiple devices. Check out [this section of the Flax scale-up guide](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-sharded-model-from-a-checkpoint) to learn how to derive a sharding tree and use it to load your checkpoint.
+In a multi-host/multi-process environment, you would want to restore your checkpoint as sharded across multiple devices. Check out the [Load sharded model from a checkpoint](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-sharded-model-from-a-checkpoint) section in the Flax [Scale up on multiple devices](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) guide to learn how to derive a sharding pytree and use it to load your checkpoint.
+
+> **Note:** JAX provides several ways to scale up your code on multiple hosts at the same time. This usually happens when the number of devices (CPU/GPU/TPU) is so large that different devices are managed by different hosts (CPU). Check out JAX’s [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html), [Using JAX in multi-host and multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html), [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html), and [Manual parallelism with `shard_map`](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html).
+
++++
 
 ## Other checkpointing features
 
-This guide only use the simplest Orbax save/load API ([`orbax.checkpoint.StandardCheckpointer`](https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html#standardcheckpointer)) to showcase all the tricks on the Flax modeling side. Feel free to use other tools or libraries as you see fit.
+This guide only uses the simplest [`orbax.checkpoint.StandardCheckpointer`](https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html#standardcheckpointer) API to show how to save and load on the Flax modeling side. Feel free to use other tools or libraries as you see fit.
 
-Check out [the Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for other commonly used features, such as:
+In addition, check out [the Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for other commonly used features, such as:
 
-* [CheckpointManager](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) to track checkpoints from different steps
+* [`CheckpointManager`](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) to track checkpoints from different steps.
 
-* [Asynchronous checkpointing](https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html)
+* [Asynchronous checkpointing](https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html).
 
-* [Transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html): a way to modify pytree structure during loading time (not after loading time, as in this guide)
+* [Orbax transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html): A way to modify pytree structure during loading time, instead of after loading time, which is demonstrated in this guide.
diff --git a/docs_nnx/guides/flax_gspmd.ipynb b/docs_nnx/guides/flax_gspmd.ipynb
index 4011c56c2..44dcfd513 100644
--- a/docs_nnx/guides/flax_gspmd.ipynb
+++ b/docs_nnx/guides/flax_gspmd.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# Scale up on multiple devices\n",
     "\n",
-    "This guide demonstrates how to scale up [Flax NNX `Module`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module) on multiple devices and hosts, such as GPUs, Google TPUs, and CPUs, using [JAX just-in-time compilation machinery (`jax.jit`)](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html)."
+    "This guide demonstrates how to scale up [Flax NNX `Module`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module) on [multiple devices and hosts](Multi-host and multi-process environments) - such as GPUs, Google TPUs, and CPUs - using the [JAX just-in-time compilation machinery (`jax.jit`)](https://jax.readthedocs.io/en/latest/jit-compilation.html) and [`flax.nnx.spmd`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html)."
    ]
   },
   {
@@ -16,9 +16,11 @@
    "source": [
     "## Overview\n",
     "\n",
-    "Flax relies on [JAX](https://jax.readthedocs.io) for numeric computations and scaling the computations up across multiple devices, such as GPU and TPUs. At the core of scaling up is the [JAX just-in-time compiler `jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html). Throughout this guide, you will be using Flax’s own [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) which wraps around [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) and works more conveniently with Flax NNX `Module`s.\n",
+    "Flax relies on [JAX](https://jax.readthedocs.io) for numeric computations and scaling the computations up across multiple devices, such as GPU and Google TPUs. At the core of scaling up is the [JAX just-in-time (`jax.jit`) compiler `jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html). Throughout this guide, you will be using Flax’s own [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) transform, which wraps around [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) and works more conveniently with Flax NNX `Module`s.\n",
     "\n",
-    "JAX compilation follows the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm. This means you write Python code as if it runs only on one device, and [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) will automatically compile and run it on multiple devices.\n",
+    "> **Note:** To learn more about Flax’s transformations, such as `nnx.jit` and `nnx.vmap`, go to [Why Flax NNX? - Transforms](https://flax.readthedocs.io/en/latest/why.html#transforms), [Transformations](https://flax.readthedocs.io/en/latest/guides/transforms.html), and [Flax NNX vs JAX Transformations](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html).\n",
+    "\n",
+    "JAX compilation follows the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm. This means you write Python code as if it runs only on one device, and [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) will [automatically compile](https://jax.readthedocs.io/en/latest/jit-compilation.html#jit-compilation) and [run it](https://jax.readthedocs.io/en/latest/sharded-computation.html) on [multiple devices](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html).\n",
     "\n",
     "To ensure the compilation performance, you often need to instruct JAX how your model's variables need to be sharded across devices. This is where Flax NNX's Sharding Metadata API - [`flax.nnx.spmd`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html) - comes in. It helps you annotate your model variables with this information.\n",
     "\n",
@@ -46,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -56,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -74,14 +76,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "We have 8 fake JAX devices now: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\n"
+      "You have 8 “fake” JAX devices now: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\n"
      ]
     }
    ],
@@ -106,7 +108,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [
     {
@@ -131,16 +133,16 @@
    "source": [
     "## Define a model with specified sharding\n",
     "\n",
-    "Next, create an example layer called `DotReluDot` that subclasses Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module). This layer carries out two dot product multiplication upon the input `x`, and uses the `jax.nn.relu` (ReLU) activation function in-between.\n",
-    "\n",
-    "To annotate a model variable with their ideal sharding, you can use [`flax.nnx.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.with_partitioning) to wrap over its initializer function. Essentially, this calls [`flax.nnx.with_metadata`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.with_metadata) which adds a `.sharding` attribute field to the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable).\n",
+    "Next, create an example layer called `DotReluDot` that subclasses Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module).\n",
+    "- This layer carries out two dot product multiplications upon the input `x`, and uses the `jax.nn.relu` (ReLU) activation function in-between.\n",
+    "- To annotate a model variable with their ideal sharding, you can use [`flax.nnx.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.with_partitioning) to wrap over its initializer function. Essentially, this calls [`flax.nnx.with_metadata`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.with_metadata) which adds a `.sharding` attribute field to the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable).\n",
     "\n",
     "> **Note:** This annotation will be [preserved and adjusted accordingly across lifted transformations in Flax NNX](https://flax.readthedocs.io/en/latest/guides/transforms.html#axes-metadata). This means if you use sharding annotations along with any transform that modifies axes (like [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html), [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html)), you need to provide sharding of that additional axis via the `transform_metadata` arg. Check out the [Flax NNX transformations (transforms) guide](https://flax.readthedocs.io/en/latest/guides/transforms.html) to learn more."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -197,12 +199,12 @@
    "source": [
     "## Initialize a sharded model\n",
     "\n",
-    "Now, you have annotations attached to the [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), but the actual weights haven't been sharded yet. If you just go ahead and create this model, all JAX arrays are still stuck in device `0`. In practice, you'd want to avoid this, because a large model will \"OOM\" (will cause the device to run out of memory) in this situation, while all the other devices are not utilized."
+    "Now, you have annotations attached to the Flax [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), but the actual weights have not been sharded yet. If you just go ahead and create this model, all [`jax.Array`s](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) are still stuck in device `0`. In practice, you'd want to avoid this, because a large model will \"OOM\" (will cause the device to run out of memory) in this situation, while all the other devices are not utilized."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
@@ -219,7 +221,7 @@
    "source": [
     "unsharded_model = DotReluDot(1024, rngs=nnx.Rngs(0))\n",
     "\n",
-    "# You have annotations sticked there, yay!\n",
+    "# You have annotations stuck there, yay!\n",
     "print(unsharded_model.dot1.kernel.sharding)     # (None, 'model')\n",
     "print(unsharded_model.w2.sharding)              # ('model', None)\n",
     "\n",
@@ -232,7 +234,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Here, you should leverage JAX's compilation mechanism, via [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit), to create the sharded model. The key is to initialize a model and assign shardings upon the model state within a `jit`ted function:\n",
+    "Here, you should leverage JAX's compilation mechanism via Flax’s [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) to create the sharded model. The key is to initialize a model and assign shardings upon the model state within a `jit`ted function:\n",
     "\n",
     "1. Use [`nnx.get_partition_spec`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_partition_spec) to strip out the `.sharding` annotations attached upon model variables.\n",
     "\n",
@@ -240,7 +242,7 @@
     "\n",
     "1. Throw away the unsharded state and return the model based upon the sharded state.\n",
     "\n",
-    "1. Compile the whole function with `nnx.jit`, which allows the output to be a stateful NNX module.\n",
+    "1. Compile the whole function with `nnx.jit`, which allows the output to be a stateful Flax NNX `Module`.\n",
     "\n",
     "1. Run it under a device mesh context so that JAX knows which devices to shard it to.\n",
     "\n",
@@ -249,7 +251,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
@@ -291,12 +293,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "You can view the sharding of any 1-D or 2-D array with `jax.debug.visualize_array_sharding`:"
+    "You can view the sharding of any 1-D or 2-D array with [`jax.debug.visualize_array_sharding`](https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.visualize_array_sharding.html):"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
     {
@@ -309,31 +311,31 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> CPU 0,4 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\"> CPU 1,5 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\"> CPU 2,6 </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\"> CPU 3,7 </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┐\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "└───────┴───────┴───────┴───────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
+       "┌───────┬───────┬───────┬───────┐\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m│CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m│CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m│CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "└───────┴───────┴───────┴───────┘\n"
       ]
      },
      "metadata": {},
@@ -349,33 +351,27 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         CPU 0,4         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         CPU 1,5         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         CPU 2,6         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         CPU 3,7         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                         </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────────┐\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>        │\n",
+       "└───────────────────────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m                         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m                         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m                         \u001b[0m\n"
+       "┌───────────────────────┐\n",
+       "│        CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m        │\n",
+       "└───────────────────────┘\n"
       ]
      },
      "metadata": {},
@@ -399,7 +395,7 @@
     "\n",
     "> **Note:** Both [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html) and [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) in the JAX documentation cover automatic parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), and semi-automatic parallelization with `jax.jit` and [`jax.lax.with_sharding_constraint](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) in greater detail.\n",
     "\n",
-    "You may have noticed you  also used [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) once in the model definition too, to constraint the sharding of an intermediate value. This is just to show that you can always use it orthogonally with the Flax NNX API, if you want to explicitly shard values that are not model variables.\n",
+    "You may have noticed you also used [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) once in the model definition to constraint the sharding of an intermediate value. This is just to show that you can always use it orthogonally with the Flax NNX API if you want to explicitly shard values that are not model variables.\n",
     "\n",
     "This brings a question: Why use the Flax NNX Annotation API then? Why not just add JAX sharding constraints inside the model definition? The most important reason is that you still need the explicit annotations to load a sharded model from an on-disk checkpoint. This is described in the next section."
    ]
@@ -408,48 +404,48 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Load sharded model from a checkpoint\n",
+    "## Load a sharded model from a checkpoint\n",
     "\n",
-    "Now you can initialize a sharded model without OOM, but what about loading it from a checkpoint on disk? JAX checkpointing libraries, such as [Orbax](https://orbax.readthedocs.io/en/latest/), usually support loading it sharded if a sharding pytree is given.\n",
+    "Now you learned how to initialize a sharded model without OOM, but what about loading it from a checkpoint on disk? JAX checkpointing libraries, such as [Orbax](https://orbax.readthedocs.io/en/latest/), usually support loading a model sharded if a sharding pytree is provided.\n",
     "\n",
-    "You can generate such as sharding pytree with [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding). To avoid any real memory allocation, use the [`nnx.eval_shape`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.eval_shape) transform to generate a model of abstract JAX arrays, and only use its `.sharding` annotations to obtain the sharding tree.\n",
+    "You can generate such a sharding pytree with Flax’s [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding). To avoid any real memory allocation, use the [`nnx.eval_shape`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.eval_shape) transform to generate a model of abstract JAX arrays, and only use its `.sharding` annotations to obtain the sharding tree.\n",
     "\n",
-    "Below is an example demonstration using Orbax's `StandardCheckpointer` API. Check out [Orbax website](https://orbax.readthedocs.io/en/latest/) to learn their latest updates and recommended APIs."
+    "Below is an example that demonstrates using Orbax's `StandardCheckpointer` API. (Go to the [Orbax documentation site](https://orbax.readthedocs.io/en/latest/) to learn about their latest and most recommended APIs.)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> CPU 0,4 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\"> CPU 1,5 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\"> CPU 2,6 </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\"> CPU 3,7 </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┐\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "└───────┴───────┴───────┴───────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
+       "┌───────┬───────┬───────┬───────┐\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m│CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m│CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m│CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "└───────┴───────┴───────┴───────┘\n"
       ]
      },
      "metadata": {},
@@ -458,33 +454,27 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         CPU 0,4         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         CPU 1,5         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         CPU 2,6         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         CPU 3,7         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                         </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────────┐\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>        │\n",
+       "└───────────────────────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m                         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m                         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m                         \u001b[0m\n"
+       "┌───────────────────────┐\n",
+       "│        CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m        │\n",
+       "└───────────────────────┘\n"
       ]
      },
      "metadata": {},
@@ -510,7 +500,7 @@
     "  abs_state, nnx.get_named_sharding(abs_state, mesh)\n",
     ")\n",
     "loaded_sharded = checkpointer.restore(path / 'checkpoint_name',\n",
-    "                                      args=ocp.args.StandardRestore(abs_state))\n",
+    "                                      target=abs_state)\n",
     "jax.debug.visualize_array_sharding(loaded_sharded.dot1.kernel.value)\n",
     "jax.debug.visualize_array_sharding(loaded_sharded.w2.value)"
    ]
@@ -521,16 +511,18 @@
    "source": [
     "## Compile the training loop\n",
     "\n",
-    "Now, from initialization or from checkpoint, you have a sharded model. To carry out the compiled, scaled up training, you need to shard the inputs as well. In this data parallelism example, the training data has its batch dimension sharded across the `data` device axis, so you should put your data in sharding `('data', None)`. You can use [`jax.device_put`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html#jax.device_put) for this.\n",
+    "Now, after either initialization or loading the checkpoint, you have a sharded model. To carry out the compiled scaled up training, you need to shard the inputs as well.\n",
     "\n",
-    "Note that with the correct sharding for all inputs, the output will be sharded in the most natural way even without JIT compilation. In the example below even without [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) on the output `y`, it was still sharded as `('data', None)`.\n",
+    "- In the data parallelism example, the training data has its batch dimension sharded across the `data` device axis, so you should put your data in sharding `('data', None)`. You can use [`jax.device_put`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html#jax.device_put) for this.\n",
+    "- Note that with the correct sharding for all inputs, the output will be sharded in the most natural way even without `jit` compilation. \n",
+    "- In the example below, even without [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) on the output `y`, it was still sharded as `('data', None)`.\n",
     "\n",
-    "> If you are interested in why: The second matmul of `DotReluDot.__call__` has two inputs of sharding `('data', 'model')` and `('model', None)`, in which both inputs' contraction axis are `model`. So a reduce-scatter matmul happened and will naturally shard the output as `('data', None)`. Check out the [JAX shard map collective guide](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html#example-2-psum-scatter-the-result) and its examples if you want to learn mathematically how it happens at a low level."
+    "> If you are interested in why: The second matmul of `DotReluDot.__call__` has two inputs of sharding `('data', 'model')` and `('model', None)`, in which both inputs' contraction axis are `model`. So a reduce-scatter matmul happened and will naturally shard the output as `('data', None)`. Check out the [JAX shard map collective guide](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html#example-2-psum-scatter-the-result) and its examples if you want to learn mathematically how it happens at a low-level."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
@@ -543,33 +535,31 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                  CPU 0,1,2,3                                   </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                  CPU 4,5,6,7                                   </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────────────────────────────────────────────────────────────────────────┐\n",
+       "│                                                                              │\n",
+       "│                                 CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                                  │\n",
+       "│                                                                              │\n",
+       "│                                                                              │\n",
+       "├──────────────────────────────────────────────────────────────────────────────┤\n",
+       "│                                                                              │\n",
+       "│                                 CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                                  │\n",
+       "│                                                                              │\n",
+       "│                                                                              │\n",
+       "└──────────────────────────────────────────────────────────────────────────────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                   \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 4,5,6,7\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                   \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n"
+       "┌──────────────────────────────────────────────────────────────────────────────┐\n",
+       "│                                                                              │\n",
+       "│                                 CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m                                  │\n",
+       "│                                                                              │\n",
+       "│                                                                              │\n",
+       "├──────────────────────────────────────────────────────────────────────────────┤\n",
+       "│                                                                              │\n",
+       "│                                 CPU \u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m                                  │\n",
+       "│                                                                              │\n",
+       "│                                                                              │\n",
+       "└──────────────────────────────────────────────────────────────────────────────┘\n"
       ]
      },
      "metadata": {},
@@ -577,28 +567,28 @@
     }
    ],
    "source": [
-    "# In data parallelism, the first dimension (batch) will be sharded on `data` axis.\n",
+    "# In data parallelism, the first dimension (batch) will be sharded on the `data` axis.\n",
     "data_sharding = NamedSharding(mesh, PartitionSpec('data', None))\n",
     "input = jax.device_put(jnp.ones((8, 1024)), data_sharding)\n",
     "\n",
     "with mesh:\n",
     "  output = sharded_model(input)\n",
     "print(output.shape)\n",
-    "jax.debug.visualize_array_sharding(output)  # Also sharded as ('data', None)"
+    "jax.debug.visualize_array_sharding(output)  # Also sharded as `('data', None)`."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now the rest of the training loop is pretty conventional - almost the same as the example in [NNX Basics](https://flax.readthedocs.io/en/latest/nnx_basics.html#transforms), except that the inputs and labels are also explicitly sharded.\n",
-    "\n",
-    "[`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) will adjust and automatically choose the best layout based on how its inputs are already sharded, so try out different shardings for your own model and inputs."
+    "Now the rest of the training loop is pretty conventional - it is almost the same as the example in [Flax NNX Basics](https://flax.readthedocs.io/en/latest/nnx_basics.html#transforms):\n",
+    "- Except that the inputs and labels are also explicitly sharded.\n",
+    "- [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) will adjust and automatically choose the best layout based on how its inputs are already sharded, so try out different shardings for your own model and inputs."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
@@ -642,19 +632,19 @@
    "source": [
     "## Profiling\n",
     "\n",
-    "If you are running on a TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance:"
+    "If you are using a Google TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 26,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "7.09 ms ± 390 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+      "7.89 ms ± 486 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
      ]
     }
    ],
@@ -678,12 +668,12 @@
     "\n",
     "JAX's [automatic](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) [SPMD]((https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD)) encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you have the option to annotate with more descriptive axis names (not just device mesh axis names like `'data'` and `'model'`), as long as you provide a mapping from your alias to the device mesh axes.\n",
     "\n",
-    "You can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot` example below."
+    "You can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot()` example below."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 27,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -723,42 +713,42 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "If you didn't provide all `sharding_rule` annotations in model definition, you can write a few lines to add it to [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) of the model, before the call of [`nnx.get_partition_spec`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_partition_spec) or [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding)."
+    "If you didn't provide all `sharding_rule` annotations in the model definition, you can write a few lines to add it to Flax’s [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) of the model, before the call of [`nnx.get_partition_spec`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_partition_spec) or [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding)."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 28,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> CPU 0,4 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\"> CPU 1,5 </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\"> CPU 2,6 </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\"> CPU 3,7 </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┐\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "└───────┴───────┴───────┴───────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
+       "┌───────┬───────┬───────┬───────┐\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m│CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m│CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m│CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "│       │       │       │       │\n",
+       "└───────┴───────┴───────┴───────┘\n"
       ]
      },
      "metadata": {},
@@ -767,33 +757,27 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         CPU 0,4         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         CPU 1,5         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         CPU 2,6         </span>\n",
-       "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         CPU 3,7         </span>\n",
-       "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                         </span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────────┐\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>        │\n",
+       "└───────────────────────┘\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;222;158;214m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m                         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\n",
-       "\u001b[38;2;255;255;255;48;2;173;73;74m                         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m                         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
-       "\u001b[38;2;0;0;0;48;2;181;207;107m                         \u001b[0m\n"
+       "┌───────────────────────┐\n",
+       "│        CPU \u001b[1;36m0\u001b[0m,\u001b[1;36m4\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m1\u001b[0m,\u001b[1;36m5\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m2\u001b[0m,\u001b[1;36m6\u001b[0m        │\n",
+       "├───────────────────────┤\n",
+       "│        CPU \u001b[1;36m3\u001b[0m,\u001b[1;36m7\u001b[0m        │\n",
+       "└───────────────────────┘\n"
       ]
      },
      "metadata": {},
@@ -850,9 +834,9 @@
     "\n",
     "  * For a simpler model, this can save you a few extra lines of code of converting the logical naming back to the device naming.\n",
     "\n",
-    "  * Shardings of intermediate *activation* values can only be done via [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) and device mesh axis. So if you want super fine-grained control over your model's sharding, directly using device mesh axis names everywhere might be less confusing.\n",
+    "  * Shardings of intermediate *activation* values can only be done via [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) and device mesh axis. Therefore, if you want super fine-grained control over your model's sharding, directly using device mesh axis names everywhere might be less confusing.\n",
     "\n",
-    "* **Logical naming**: Helpful if you want to experiment around and find the most optimal partition layout for your *model weights*."
+    "* **Logical naming**: This is helpful if you want to experiment around and find the most optimal partition layout for your *model weights*."
    ]
   }
  ],
@@ -870,7 +854,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,
diff --git a/docs_nnx/guides/flax_gspmd.md b/docs_nnx/guides/flax_gspmd.md
index e13c4f2c4..50441f941 100644
--- a/docs_nnx/guides/flax_gspmd.md
+++ b/docs_nnx/guides/flax_gspmd.md
@@ -10,15 +10,17 @@ jupytext:
 
 # Scale up on multiple devices
 
-This guide demonstrates how to scale up [Flax NNX `Module`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module) on multiple devices and hosts, such as GPUs, Google TPUs, and CPUs, using [JAX just-in-time compilation machinery (`jax.jit`)](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html).
+This guide demonstrates how to scale up [Flax NNX `Module`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module) on [multiple devices and hosts](Multi-host and multi-process environments) - such as GPUs, Google TPUs, and CPUs - using the [JAX just-in-time compilation machinery (`jax.jit`)](https://jax.readthedocs.io/en/latest/jit-compilation.html) and [`flax.nnx.spmd`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html).
 
 +++
 
 ## Overview
 
-Flax relies on [JAX](https://jax.readthedocs.io) for numeric computations and scaling the computations up across multiple devices, such as GPU and TPUs. At the core of scaling up is the [JAX just-in-time compiler `jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html). Throughout this guide, you will be using Flax’s own [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) which wraps around [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) and works more conveniently with Flax NNX `Module`s.
+Flax relies on [JAX](https://jax.readthedocs.io) for numeric computations and scaling the computations up across multiple devices, such as GPU and Google TPUs. At the core of scaling up is the [JAX just-in-time (`jax.jit`) compiler `jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html). Throughout this guide, you will be using Flax’s own [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) transform, which wraps around [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) and works more conveniently with Flax NNX `Module`s.
 
-JAX compilation follows the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm. This means you write Python code as if it runs only on one device, and [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) will automatically compile and run it on multiple devices.
+> **Note:** To learn more about Flax’s transformations, such as `nnx.jit` and `nnx.vmap`, go to [Why Flax NNX? - Transforms](https://flax.readthedocs.io/en/latest/why.html#transforms), [Transformations](https://flax.readthedocs.io/en/latest/guides/transforms.html), and [Flax NNX vs JAX Transformations](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html).
+
+JAX compilation follows the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm. This means you write Python code as if it runs only on one device, and [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) will [automatically compile](https://jax.readthedocs.io/en/latest/jit-compilation.html#jit-compilation) and [run it](https://jax.readthedocs.io/en/latest/sharded-computation.html) on [multiple devices](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html).
 
 To ensure the compilation performance, you often need to instruct JAX how your model's variables need to be sharded across devices. This is where Flax NNX's Sharding Metadata API - [`flax.nnx.spmd`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html) - comes in. It helps you annotate your model variables with this information.
 
@@ -79,9 +81,9 @@ print(mesh)
 
 ## Define a model with specified sharding
 
-Next, create an example layer called `DotReluDot` that subclasses Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module). This layer carries out two dot product multiplication upon the input `x`, and uses the `jax.nn.relu` (ReLU) activation function in-between.
-
-To annotate a model variable with their ideal sharding, you can use [`flax.nnx.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.with_partitioning) to wrap over its initializer function. Essentially, this calls [`flax.nnx.with_metadata`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.with_metadata) which adds a `.sharding` attribute field to the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable).
+Next, create an example layer called `DotReluDot` that subclasses Flax [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html#flax.nnx.Module).
+- This layer carries out two dot product multiplications upon the input `x`, and uses the `jax.nn.relu` (ReLU) activation function in-between.
+- To annotate a model variable with their ideal sharding, you can use [`flax.nnx.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.with_partitioning) to wrap over its initializer function. Essentially, this calls [`flax.nnx.with_metadata`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.with_metadata) which adds a `.sharding` attribute field to the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable).
 
 > **Note:** This annotation will be [preserved and adjusted accordingly across lifted transformations in Flax NNX](https://flax.readthedocs.io/en/latest/guides/transforms.html#axes-metadata). This means if you use sharding annotations along with any transform that modifies axes (like [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html), [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html)), you need to provide sharding of that additional axis via the `transform_metadata` arg. Check out the [Flax NNX transformations (transforms) guide](https://flax.readthedocs.io/en/latest/guides/transforms.html) to learn more.
 
@@ -130,12 +132,12 @@ JAX's [Distributed arrays and automatic parallelization](https://jax.readthedocs
 
 ## Initialize a sharded model
 
-Now, you have annotations attached to the [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), but the actual weights haven't been sharded yet. If you just go ahead and create this model, all JAX arrays are still stuck in device `0`. In practice, you'd want to avoid this, because a large model will "OOM" (will cause the device to run out of memory) in this situation, while all the other devices are not utilized.
+Now, you have annotations attached to the Flax [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), but the actual weights have not been sharded yet. If you just go ahead and create this model, all [`jax.Array`s](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) are still stuck in device `0`. In practice, you'd want to avoid this, because a large model will "OOM" (will cause the device to run out of memory) in this situation, while all the other devices are not utilized.
 
 ```{code-cell} ipython3
 unsharded_model = DotReluDot(1024, rngs=nnx.Rngs(0))
 
-# You have annotations sticked there, yay!
+# You have annotations stuck there, yay!
 print(unsharded_model.dot1.kernel.sharding)     # (None, 'model')
 print(unsharded_model.w2.sharding)              # ('model', None)
 
@@ -144,7 +146,7 @@ print(unsharded_model.dot1.kernel.value.sharding)  # SingleDeviceSharding
 print(unsharded_model.w2.value.sharding)           # SingleDeviceSharding
 ```
 
-Here, you should leverage JAX's compilation mechanism, via [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit), to create the sharded model. The key is to initialize a model and assign shardings upon the model state within a `jit`ted function:
+Here, you should leverage JAX's compilation mechanism via Flax’s [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) to create the sharded model. The key is to initialize a model and assign shardings upon the model state within a `jit`ted function:
 
 1. Use [`nnx.get_partition_spec`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_partition_spec) to strip out the `.sharding` annotations attached upon model variables.
 
@@ -152,7 +154,7 @@ Here, you should leverage JAX's compilation mechanism, via [`nnx.jit`](https://f
 
 1. Throw away the unsharded state and return the model based upon the sharded state.
 
-1. Compile the whole function with `nnx.jit`, which allows the output to be a stateful NNX module.
+1. Compile the whole function with `nnx.jit`, which allows the output to be a stateful Flax NNX `Module`.
 
 1. Run it under a device mesh context so that JAX knows which devices to shard it to.
 
@@ -184,7 +186,7 @@ assert sharded_model.w2.value.sharding.is_equivalent_to(
 )
 ```
 
-You can view the sharding of any 1-D or 2-D array with `jax.debug.visualize_array_sharding`:
+You can view the sharding of any 1-D or 2-D array with [`jax.debug.visualize_array_sharding`](https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.visualize_array_sharding.html):
 
 ```{code-cell} ipython3
 print("sharded_model.dot1.kernel (None, 'model') :")
@@ -199,19 +201,19 @@ The key to shard a JAX array is to call [`jax.lax.with_sharding_constraint`](htt
 
 > **Note:** Both [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html) and [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) in the JAX documentation cover automatic parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), and semi-automatic parallelization with `jax.jit` and [`jax.lax.with_sharding_constraint](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) in greater detail.
 
-You may have noticed you  also used [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) once in the model definition too, to constraint the sharding of an intermediate value. This is just to show that you can always use it orthogonally with the Flax NNX API, if you want to explicitly shard values that are not model variables.
+You may have noticed you also used [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) once in the model definition to constraint the sharding of an intermediate value. This is just to show that you can always use it orthogonally with the Flax NNX API if you want to explicitly shard values that are not model variables.
 
 This brings a question: Why use the Flax NNX Annotation API then? Why not just add JAX sharding constraints inside the model definition? The most important reason is that you still need the explicit annotations to load a sharded model from an on-disk checkpoint. This is described in the next section.
 
 +++
 
-## Load sharded model from a checkpoint
+## Load a sharded model from a checkpoint
 
-Now you can initialize a sharded model without OOM, but what about loading it from a checkpoint on disk? JAX checkpointing libraries, such as [Orbax](https://orbax.readthedocs.io/en/latest/), usually support loading it sharded if a sharding pytree is given.
+Now you learned how to initialize a sharded model without OOM, but what about loading it from a checkpoint on disk? JAX checkpointing libraries, such as [Orbax](https://orbax.readthedocs.io/en/latest/), usually support loading a model sharded if a sharding pytree is provided.
 
-You can generate such as sharding pytree with [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding). To avoid any real memory allocation, use the [`nnx.eval_shape`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.eval_shape) transform to generate a model of abstract JAX arrays, and only use its `.sharding` annotations to obtain the sharding tree.
+You can generate such a sharding pytree with Flax’s [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding). To avoid any real memory allocation, use the [`nnx.eval_shape`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.eval_shape) transform to generate a model of abstract JAX arrays, and only use its `.sharding` annotations to obtain the sharding tree.
 
-Below is an example demonstration using Orbax's `StandardCheckpointer` API. Check out [Orbax website](https://orbax.readthedocs.io/en/latest/) to learn their latest updates and recommended APIs.
+Below is an example that demonstrates using Orbax's `StandardCheckpointer` API. (Go to the [Orbax documentation site](https://orbax.readthedocs.io/en/latest/) to learn about their latest and most recommended APIs.)
 
 ```{code-cell} ipython3
 import orbax.checkpoint as ocp
@@ -232,33 +234,35 @@ abs_state = jax.tree.map(
   abs_state, nnx.get_named_sharding(abs_state, mesh)
 )
 loaded_sharded = checkpointer.restore(path / 'checkpoint_name',
-                                      args=ocp.args.StandardRestore(abs_state))
+                                      target=abs_state)
 jax.debug.visualize_array_sharding(loaded_sharded.dot1.kernel.value)
 jax.debug.visualize_array_sharding(loaded_sharded.w2.value)
 ```
 
 ## Compile the training loop
 
-Now, from initialization or from checkpoint, you have a sharded model. To carry out the compiled, scaled up training, you need to shard the inputs as well. In this data parallelism example, the training data has its batch dimension sharded across the `data` device axis, so you should put your data in sharding `('data', None)`. You can use [`jax.device_put`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html#jax.device_put) for this.
+Now, after either initialization or loading the checkpoint, you have a sharded model. To carry out the compiled scaled up training, you need to shard the inputs as well.
 
-Note that with the correct sharding for all inputs, the output will be sharded in the most natural way even without JIT compilation. In the example below even without [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) on the output `y`, it was still sharded as `('data', None)`.
+- In the data parallelism example, the training data has its batch dimension sharded across the `data` device axis, so you should put your data in sharding `('data', None)`. You can use [`jax.device_put`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html#jax.device_put) for this.
+- Note that with the correct sharding for all inputs, the output will be sharded in the most natural way even without `jit` compilation. 
+- In the example below, even without [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) on the output `y`, it was still sharded as `('data', None)`.
 
-> If you are interested in why: The second matmul of `DotReluDot.__call__` has two inputs of sharding `('data', 'model')` and `('model', None)`, in which both inputs' contraction axis are `model`. So a reduce-scatter matmul happened and will naturally shard the output as `('data', None)`. Check out the [JAX shard map collective guide](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html#example-2-psum-scatter-the-result) and its examples if you want to learn mathematically how it happens at a low level.
+> If you are interested in why: The second matmul of `DotReluDot.__call__` has two inputs of sharding `('data', 'model')` and `('model', None)`, in which both inputs' contraction axis are `model`. So a reduce-scatter matmul happened and will naturally shard the output as `('data', None)`. Check out the [JAX shard map collective guide](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html#example-2-psum-scatter-the-result) and its examples if you want to learn mathematically how it happens at a low-level.
 
 ```{code-cell} ipython3
-# In data parallelism, the first dimension (batch) will be sharded on `data` axis.
+# In data parallelism, the first dimension (batch) will be sharded on the `data` axis.
 data_sharding = NamedSharding(mesh, PartitionSpec('data', None))
 input = jax.device_put(jnp.ones((8, 1024)), data_sharding)
 
 with mesh:
   output = sharded_model(input)
 print(output.shape)
-jax.debug.visualize_array_sharding(output)  # Also sharded as ('data', None)
+jax.debug.visualize_array_sharding(output)  # Also sharded as `('data', None)`.
 ```
 
-Now the rest of the training loop is pretty conventional - almost the same as the example in [NNX Basics](https://flax.readthedocs.io/en/latest/nnx_basics.html#transforms), except that the inputs and labels are also explicitly sharded.
-
-[`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) will adjust and automatically choose the best layout based on how its inputs are already sharded, so try out different shardings for your own model and inputs.
+Now the rest of the training loop is pretty conventional - it is almost the same as the example in [Flax NNX Basics](https://flax.readthedocs.io/en/latest/nnx_basics.html#transforms):
+- Except that the inputs and labels are also explicitly sharded.
+- [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) will adjust and automatically choose the best layout based on how its inputs are already sharded, so try out different shardings for your own model and inputs.
 
 ```{code-cell} ipython3
 optimizer = nnx.Optimizer(sharded_model, optax.adam(1e-3))  # reference sharing
@@ -285,7 +289,7 @@ with mesh:
 
 ## Profiling
 
-If you are running on a TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance:
+If you are using a Google TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance:
 
 ```{code-cell} ipython3
 %%timeit
@@ -302,7 +306,7 @@ with mesh:
 
 JAX's [automatic](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) [SPMD]((https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD)) encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you have the option to annotate with more descriptive axis names (not just device mesh axis names like `'data'` and `'model'`), as long as you provide a mapping from your alias to the device mesh axes.
 
-You can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot` example below.
+You can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot()` example below.
 
 ```{code-cell} ipython3
 # The mapping from alias annotation to the device mesh.
@@ -337,7 +341,7 @@ class LogicalDotReluDot(nnx.Module):
     return z
 ```
 
-If you didn't provide all `sharding_rule` annotations in model definition, you can write a few lines to add it to [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) of the model, before the call of [`nnx.get_partition_spec`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_partition_spec) or [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding).
+If you didn't provide all `sharding_rule` annotations in the model definition, you can write a few lines to add it to Flax’s [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) of the model, before the call of [`nnx.get_partition_spec`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_partition_spec) or [`nnx.get_named_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.get_named_sharding).
 
 ```{code-cell} ipython3
 def add_sharding_rule(vs: nnx.VariableState) -> nnx.VariableState:
@@ -384,6 +388,6 @@ Choosing when to use a device or logical axis depends on how much you want to co
 
   * For a simpler model, this can save you a few extra lines of code of converting the logical naming back to the device naming.
 
-  * Shardings of intermediate *activation* values can only be done via [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) and device mesh axis. So if you want super fine-grained control over your model's sharding, directly using device mesh axis names everywhere might be less confusing.
+  * Shardings of intermediate *activation* values can only be done via [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) and device mesh axis. Therefore, if you want super fine-grained control over your model's sharding, directly using device mesh axis names everywhere might be less confusing.
 
-* **Logical naming**: Helpful if you want to experiment around and find the most optimal partition layout for your *model weights*.
+* **Logical naming**: This is helpful if you want to experiment around and find the most optimal partition layout for your *model weights*.
diff --git a/flax/nnx/__init__.py b/flax/nnx/__init__.py
index 367fe7e09..f6ef81a9a 100644
--- a/flax/nnx/__init__.py
+++ b/flax/nnx/__init__.py
@@ -86,6 +86,13 @@
 from .nn.attention import dot_product_attention as dot_product_attention
 from .nn.attention import make_attention_mask as make_attention_mask
 from .nn.attention import make_causal_mask as make_causal_mask
+from .nn.recurrent import RNNCellBase as RNNCellBase
+from .nn.recurrent import LSTMCell as LSTMCell
+from .nn.recurrent import GRUCell as GRUCell
+from .nn.recurrent import OptimizedLSTMCell as OptimizedLSTMCell
+from .nn.recurrent import SimpleCell as SimpleCell
+from .nn.recurrent import RNN as RNN
+from .nn.recurrent import Bidirectional as Bidirectional
 from .nn.linear import Conv as Conv
 from .nn.linear import ConvTranspose as ConvTranspose
 from .nn.linear import Embed as Embed
diff --git a/flax/nnx/extract.py b/flax/nnx/extract.py
index 948d8f9fc..191a0c195 100644
--- a/flax/nnx/extract.py
+++ b/flax/nnx/extract.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import abc
 import contextlib
 import dataclasses
 import threading
@@ -23,7 +24,7 @@
 from flax import struct
 from flax.nnx.object import Object
 from flax.typing import Missing, PathParts
-from flax.nnx import graph
+from flax.nnx import graph, variablelib
 
 
 A = tp.TypeVar('A')
@@ -119,6 +120,14 @@ def _maybe_insert(x):
     _maybe_insert, pytree, is_leaf=lambda x: isinstance(x, ExtractionIndex)
   )
 
+class PrefixMapping(abc.ABC):
+  @abc.abstractmethod
+  def map_prefix(
+    self,
+    path: variablelib.PathParts,
+    variable: variablelib.Variable,
+    /,
+  ) -> tp.Any: ...
 
 def check_consistent_aliasing(
   node: tuple[tp.Any, ...],
@@ -143,11 +152,16 @@ def check_consistent_aliasing(
           raise ValueError(
             f'Cannot extract graph node from different trace level, got {value!r}'
           )
-      if value in node_prefixes:
-        paths_prefixes = node_prefixes[value]
-        paths_prefixes.append((path, prefix))
-      else:
-        node_prefixes[value] = [(path, prefix)]
+        if isinstance(prefix, PrefixMapping):
+          variable_prefix = prefix.map_prefix(path, value)
+        else:
+          variable_prefix = prefix
+
+        if value in node_prefixes:
+          paths_prefixes = node_prefixes[value]
+          paths_prefixes.append((path, variable_prefix))
+        else:
+          node_prefixes[value] = [(path, variable_prefix)]
 
   # check for inconsistent aliasing
   node_msgs = []
diff --git a/flax/nnx/filterlib.py b/flax/nnx/filterlib.py
index 9ad3419df..e702966cf 100644
--- a/flax/nnx/filterlib.py
+++ b/flax/nnx/filterlib.py
@@ -87,6 +87,24 @@ def __repr__(self):
     return f'PathContains({self.key!r})'
 
 
+class PathIn:
+  def __init__(self, *paths: PathParts):
+    self.paths = frozenset(paths)
+
+  def __call__(self, path: PathParts, x: tp.Any):
+    return path in self.paths
+
+  def __repr__(self):
+    paths_repr = ','.join(map(repr, self.paths))
+    return f'PathIn({paths_repr})'
+
+  def __eq__(self, other):
+    return isinstance(other, PathIn) and self.paths == other.paths
+
+  def __hash__(self):
+    return hash(self.paths)
+
+
 @dataclasses.dataclass(frozen=True)
 class OfType:
   type: type
diff --git a/flax/nnx/nn/recurrent.py b/flax/nnx/nn/recurrent.py
new file mode 100644
index 000000000..e659144af
--- /dev/null
+++ b/flax/nnx/nn/recurrent.py
@@ -0,0 +1,923 @@
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""RNN modules for Flax."""
+
+from typing import (
+    Any,
+    TypeVar
+)
+from collections.abc import Callable
+from functools import partial
+from typing_extensions import Protocol
+from absl import logging
+
+import jax
+import jax.numpy as jnp
+
+from flax import nnx
+from flax.nnx import rnglib
+from flax.nnx.module import Module
+from flax.nnx.nn import initializers
+from flax.nnx.nn.linear import Linear
+from flax.nnx.nn.activations import sigmoid
+from flax.nnx.nn.activations import tanh
+from flax.nnx.transforms.iteration import Carry
+from flax.typing import (
+    Dtype,
+    Initializer,
+    Shape
+)
+
+default_kernel_init = initializers.lecun_normal()
+default_bias_init = initializers.zeros_init()
+
+A = TypeVar("A")
+Array = jax.Array
+Output = Any
+
+
+class RNNCellBase(Module):
+    """RNN cell base class."""
+
+    def initialize_carry(
+        self, input_shape: tuple[int, ...], rngs: rnglib.Rngs | None = None
+    ) -> Carry:
+        """Initialize the RNN cell carry.
+
+        Args:
+          rng: random number generator passed to the init_fn.
+          input_shape: a tuple providing the shape of the input to the cell.
+
+        Returns:
+          An initialized carry for the given RNN cell.
+        """
+        raise NotImplementedError
+
+    def __call__(
+        self,
+        carry: Carry,
+        inputs: Array
+    ) -> tuple[Carry, Array]:
+        """Run the RNN cell.
+
+        Args:
+          carry: the hidden state of the RNN cell.
+          inputs: an ndarray with the input for the current time step.
+            All dimensions except the final are considered batch dimensions.
+
+        Returns:
+          A tuple with the new carry and the output.
+        """
+        raise NotImplementedError
+
+    @property
+    def num_feature_axes(self) -> int:
+        """Returns the number of feature axes of the RNN cell."""
+        raise NotImplementedError
+
+def modified_orthogonal(key: Array, shape: Shape, dtype: Dtype = jnp.float32) -> Array:
+    """Modified orthogonal initializer for compatibility with half precision."""
+    initializer = initializers.orthogonal()
+    return initializer(key, shape).astype(dtype)
+
+class LSTMCell(RNNCellBase):
+    r"""LSTM cell.
+
+  The mathematical definition of the cell is as follows
+
+  .. math::
+      \begin{array}{ll}
+      i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\
+      f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\
+      g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\
+      o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\
+      c' = f * c + i * g \\
+      h' = o * \tanh(c') \\
+      \end{array}
+
+  where x is the input, h is the output of the previous time step, and c is
+  the memory.
+  """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int,
+        *,
+        gate_fn: Callable[..., Any] = sigmoid,
+        activation_fn: Callable[..., Any] = tanh,
+        kernel_init: Initializer = default_kernel_init,
+        recurrent_kernel_init: Initializer = modified_orthogonal,
+        bias_init: Initializer = initializers.zeros_init(),
+        dtype: Dtype | None = None,
+        param_dtype: Dtype = jnp.float32,
+        carry_init: Initializer = initializers.zeros_init(),
+        rngs: rnglib.Rngs,
+    ):
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.gate_fn = gate_fn
+        self.activation_fn = activation_fn
+        self.kernel_init = kernel_init
+        self.recurrent_kernel_init = recurrent_kernel_init
+        self.bias_init = bias_init
+        self.dtype = dtype
+        self.param_dtype = param_dtype
+        self.carry_init = carry_init
+        self.rngs = rngs
+
+        # input and recurrent layers are summed so only one needs a bias.
+        dense_i = partial(
+            Linear,
+            in_features=in_features,
+            out_features=hidden_features,
+            use_bias=False,
+            kernel_init=self.kernel_init,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            rngs=rngs,
+        )
+
+        dense_h = partial(
+            Linear,
+            in_features=hidden_features,
+            out_features=hidden_features,
+            use_bias=True,
+            kernel_init=self.recurrent_kernel_init,
+            bias_init=self.bias_init,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            rngs=rngs,
+        )
+
+        self.ii = dense_i()
+        self.if_ = dense_i()
+        self.ig = dense_i()
+        self.io = dense_i()
+        self.hi = dense_h()
+        self.hf = dense_h()
+        self.hg = dense_h()
+        self.ho = dense_h()
+
+    def __call__(self, carry: tuple[Array, Array], inputs: Array) -> tuple[tuple[Array, Array], Array]: # type: ignore[override]
+        r"""A long short-term memory (LSTM) cell.
+
+        Args:
+          carry: the hidden state of the LSTM cell,
+            initialized using ``LSTMCell.initialize_carry``.
+          inputs: an ndarray with the input for the current time step.
+            All dimensions except the final are considered batch dimensions.
+
+        Returns:
+          A tuple with the new carry and the output.
+        """
+        c, h = carry
+        i = self.gate_fn(self.ii(inputs) + self.hi(h))
+        f = self.gate_fn(self.if_(inputs) + self.hf(h))
+        g = self.activation_fn(self.ig(inputs) + self.hg(h))
+        o = self.gate_fn(self.io(inputs) + self.ho(h))
+        new_c = f * c + i * g
+        new_h = o * self.activation_fn(new_c)
+        return (new_c, new_h), new_h
+
+    def initialize_carry(self, input_shape: tuple[int, ...], rngs: rnglib.Rngs | None = None) -> tuple[Array, Array]: # type: ignore[override]
+        """Initialize the RNN cell carry.
+
+        Args:
+          rng: random number generator passed to the init_fn.
+          input_shape: a tuple providing the shape of the input to the cell.
+        Returns:
+          An initialized carry for the given RNN cell.
+        """
+        batch_dims = input_shape[:-1]
+        if rngs is None:
+            rngs = self.rngs
+        mem_shape = batch_dims + (self.hidden_features,)
+        c = self.carry_init(rngs.carry(), mem_shape, self.param_dtype)
+        h = self.carry_init(rngs.carry(), mem_shape, self.param_dtype)
+        return (c, h)
+
+    @property
+    def num_feature_axes(self) -> int:
+        return 1
+
+
+class OptimizedLSTMCell(RNNCellBase):
+    r"""More efficient LSTM Cell that concatenates state components before matmul.
+
+    The parameters are compatible with ``LSTMCell``. Note that this cell is often
+    faster than ``LSTMCell`` as long as the hidden size is roughly <= 2048 units.
+
+    The mathematical definition of the cell is the same as ``LSTMCell`` and as
+    follows:
+
+    .. math::
+
+        \begin{array}{ll}
+        i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\
+        f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\
+        g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\
+        o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\
+        c' = f * c + i * g \\
+        h' = o * \tanh(c') \\
+        \end{array}
+
+    where x is the input, h is the output of the previous time step, and c is
+    the memory.
+
+    Attributes:
+        gate_fn: activation function used for gates (default: sigmoid).
+        activation_fn: activation function used for output and memory update
+          (default: tanh).
+        kernel_init: initializer function for the kernels that transform
+          the input (default: lecun_normal).
+        recurrent_kernel_init: initializer function for the kernels that transform
+          the hidden state (default: initializers.orthogonal()).
+        bias_init: initializer for the bias parameters (default: initializers.zeros_init()).
+        dtype: the dtype of the computation (default: infer from inputs and params).
+        param_dtype: the dtype passed to parameter initializers (default: float32).
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int,
+        *,
+        gate_fn: Callable[..., Any] = sigmoid,
+        activation_fn: Callable[..., Any] = tanh,
+        kernel_init: Initializer = default_kernel_init,
+        recurrent_kernel_init: Initializer = initializers.orthogonal(),
+        bias_init: Initializer = initializers.zeros_init(),
+        dtype: Dtype | None = None,
+        param_dtype: Dtype = jnp.float32,
+        carry_init: Initializer = initializers.zeros_init(),
+        rngs: rnglib.Rngs,
+    ):
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.gate_fn = gate_fn
+        self.activation_fn = activation_fn
+        self.kernel_init = kernel_init
+        self.recurrent_kernel_init = recurrent_kernel_init
+        self.bias_init = bias_init
+        self.dtype = dtype
+        self.param_dtype = param_dtype
+        self.carry_init = carry_init
+        self.rngs = rngs
+
+        # input and recurrent layers are summed so only one needs a bias.
+        self.dense_i = Linear(
+            in_features=in_features,
+            out_features=4 * hidden_features,
+            use_bias=False,
+            kernel_init=self.kernel_init,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            rngs=rngs,
+        )
+
+        self.dense_h = Linear(
+            in_features=hidden_features,
+            out_features=4 * hidden_features,
+            use_bias=True,
+            kernel_init=self.recurrent_kernel_init,
+            bias_init=self.bias_init,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            rngs=rngs,
+        )
+
+    def __call__(self, carry: tuple[Array, Array], inputs: Array) -> tuple[tuple[Array, Array], Array]: # type: ignore[override]
+        r"""An optimized long short-term memory (LSTM) cell.
+
+        Args:
+          carry: the hidden state of the LSTM cell, initialized using
+            ``LSTMCell.initialize_carry``.
+          inputs: an ndarray with the input for the current time step.
+            All dimensions except the final are considered batch dimensions.
+
+        Returns:
+          A tuple with the new carry and the output.
+        """
+        c, h = carry
+
+        # Compute combined transformations for inputs and hidden state
+        y = self.dense_i(inputs) + self.dense_h(h)
+
+        # Split the combined transformations into individual gates
+        i, f, g, o = jnp.split(y, indices_or_sections=4, axis=-1)
+
+        # Apply gate activations
+        i = self.gate_fn(i)
+        f = self.gate_fn(f)
+        g = self.activation_fn(g)
+        o = self.gate_fn(o)
+
+        # Update cell state and hidden state
+        new_c = f * c + i * g
+        new_h = o * self.activation_fn(new_c)
+        return (new_c, new_h), new_h
+
+    def initialize_carry(self, input_shape: tuple[int, ...], rngs: rnglib.Rngs | None = None) -> tuple[Array, Array]: # type: ignore[override]
+        """Initialize the RNN cell carry.
+
+        Args:
+          rngs: random number generator passed to the init_fn.
+          input_shape: a tuple providing the shape of the input to the cell.
+
+        Returns:
+          An initialized carry for the given RNN cell.
+        """
+        batch_dims = input_shape[:-1]
+        if rngs is None:
+            rngs = self.rngs
+        mem_shape = batch_dims + (self.hidden_features,)
+        c = self.carry_init(rngs.carry(), mem_shape, self.param_dtype)
+        h = self.carry_init(rngs.carry(), mem_shape, self.param_dtype)
+        return (c, h)
+
+    @property
+    def num_feature_axes(self) -> int:
+        return 1
+
+
+class SimpleCell(RNNCellBase):
+    r"""Simple cell.
+
+    The mathematical definition of the cell is as follows
+
+    .. math::
+
+        \begin{array}{ll}
+        h' = \tanh(W_i x + b_i + W_h h)
+        \end{array}
+
+    where x is the input and h is the output of the previous time step.
+
+    If `residual` is `True`,
+
+    .. math::
+
+        \begin{array}{ll}
+        h' = \tanh(W_i x + b_i + W_h h + h)
+        \end{array}
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int,  # not inferred from carry for now
+        *,
+        dtype: Dtype = jnp.float32,
+        param_dtype: Dtype = jnp.float32,
+        carry_init: Initializer = initializers.zeros_init(),
+        residual: bool = False,
+        activation_fn: Callable[..., Any] = tanh,
+        kernel_init: Initializer = initializers.lecun_normal(),
+        recurrent_kernel_init: Initializer = initializers.orthogonal(),
+        bias_init: Initializer = initializers.zeros_init(),
+        rngs: rnglib.Rngs,
+    ):
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.dtype = dtype
+        self.param_dtype = param_dtype
+        self.carry_init = carry_init
+        self.residual = residual
+        self.activation_fn = activation_fn
+        self.kernel_init = kernel_init
+        self.recurrent_kernel_init = recurrent_kernel_init
+        self.bias_init = bias_init
+        self.rngs = rngs
+
+        # self.hidden_features = carry.shape[-1]
+        # input and recurrent layers are summed so only one needs a bias.
+        self.dense_h = Linear(
+            in_features=self.hidden_features,
+            out_features=self.hidden_features,
+            use_bias=False,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            kernel_init=self.recurrent_kernel_init,
+            rngs=rngs,
+        )
+        self.dense_i = Linear(
+            in_features=self.in_features,
+            out_features=self.hidden_features,
+            use_bias=True,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            kernel_init=self.kernel_init,
+            bias_init=self.bias_init,
+            rngs=rngs,
+        )
+
+    def __call__(self, carry: Array, inputs: Array) -> tuple[Array, Array]: # type: ignore[override]
+        new_carry = self.dense_i(inputs) + self.dense_h(carry)
+        if self.residual:
+            new_carry += carry
+        new_carry = self.activation_fn(new_carry)
+        return new_carry, new_carry
+
+    def initialize_carry(self, input_shape: tuple[int, ...], rngs: rnglib.Rngs | None = None) -> Array: # type: ignore[override]
+        """Initialize the RNN cell carry.
+
+        Args:
+          rng: random number generator passed to the init_fn.
+          input_shape: a tuple providing the shape of the input to the cell.
+
+        Returns:
+          An initialized carry for the given RNN cell.
+        """
+        if rngs is None:
+            rngs = self.rngs
+        batch_dims = input_shape[:-1]
+        mem_shape = batch_dims + (self.hidden_features,)
+        return self.carry_init(rngs.carry(), mem_shape, self.param_dtype)
+
+    @property
+    def num_feature_axes(self) -> int:
+        return 1
+
+
+class GRUCell(RNNCellBase):
+    r"""GRU cell.
+
+    The mathematical definition of the cell is as follows
+
+    .. math::
+
+        \begin{array}{ll}
+        r = \sigma(W_{ir} x + b_{ir} + W_{hr} h) \\
+        z = \sigma(W_{iz} x + b_{iz} + W_{hz} h) \\
+        n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\
+        h' = (1 - z) * n + z * h \\
+        \end{array}
+
+    where x is the input and h is the output of the previous time step.
+
+    Attributes:
+        in_features: number of input features.
+        hidden_features: number of output features.
+        gate_fn: activation function used for gates (default: sigmoid).
+        activation_fn: activation function used for output and memory update
+          (default: tanh).
+        kernel_init: initializer function for the kernels that transform
+          the input (default: lecun_normal).
+        recurrent_kernel_init: initializer function for the kernels that transform
+          the hidden state (default: initializers.orthogonal()).
+        bias_init: initializer for the bias parameters (default: initializers.zeros_init()).
+        dtype: the dtype of the computation (default: None).
+        param_dtype: the dtype passed to parameter initializers (default: float32).
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int,
+        *,
+        gate_fn: Callable[..., Any] = sigmoid,
+        activation_fn: Callable[..., Any] = tanh,
+        kernel_init: Initializer = default_kernel_init,
+        recurrent_kernel_init: Initializer = initializers.orthogonal(),
+        bias_init: Initializer = initializers.zeros_init(),
+        dtype: Dtype | None = None,
+        param_dtype: Dtype = jnp.float32,
+        carry_init: Initializer = initializers.zeros_init(),
+        rngs: rnglib.Rngs,
+    ):
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.gate_fn = gate_fn
+        self.activation_fn = activation_fn
+        self.kernel_init = kernel_init
+        self.recurrent_kernel_init = recurrent_kernel_init
+        self.bias_init = bias_init
+        self.dtype = dtype
+        self.param_dtype = param_dtype
+        self.carry_init = carry_init
+        self.rngs = rngs
+
+        # Combine input transformations into a single linear layer
+        self.dense_i = Linear(
+            in_features=in_features,
+            out_features=3 * hidden_features,  # r, z, n
+            use_bias=True,
+            kernel_init=self.kernel_init,
+            bias_init=self.bias_init,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            rngs=rngs,
+        )
+
+        self.dense_h = Linear(
+            in_features=hidden_features,
+            out_features=3 * hidden_features,  # r, z, n
+            use_bias=False,
+            kernel_init=self.recurrent_kernel_init,
+            dtype=self.dtype,
+            param_dtype=self.param_dtype,
+            rngs=rngs,
+        )
+
+    def __call__(self, carry: Array, inputs: Array) -> tuple[Array, Array]: # type: ignore[override]
+        """Gated recurrent unit (GRU) cell.
+
+        Args:
+            carry: the hidden state of the GRU cell,
+              initialized using ``GRUCell.initialize_carry``.
+            inputs: an ndarray with the input for the current time step.
+              All dimensions except the final are considered batch dimensions.
+
+        Returns:
+            A tuple with the new carry and the output.
+        """
+        h = carry
+
+        # Compute combined transformations for inputs and hidden state
+        x_transformed = self.dense_i(inputs)
+        h_transformed = self.dense_h(h)
+
+        # Split the combined transformations into individual components
+        xi_r, xi_z, xi_n = jnp.split(x_transformed, 3, axis=-1)
+        hh_r, hh_z, hh_n = jnp.split(h_transformed, 3, axis=-1)
+
+        # Compute gates
+        r = self.gate_fn(xi_r + hh_r)
+        z = self.gate_fn(xi_z + hh_z)
+
+        # Compute n with an additional linear transformation on h
+        n = self.activation_fn(xi_n + r * hh_n)
+
+        # Update hidden state
+        new_h = (1.0 - z) * n + z * h
+        return new_h, new_h
+
+    def initialize_carry(self, input_shape: tuple[int, ...], rngs: rnglib.Rngs | None = None) -> Array: # type: ignore[override]
+        """Initialize the RNN cell carry.
+
+        Args:
+            rngs: random number generator passed to the init_fn.
+            input_shape: a tuple providing the shape of the input to the cell.
+
+        Returns:
+            An initialized carry for the given RNN cell.
+        """
+        batch_dims = input_shape[:-1]
+        if rngs is None:
+            rngs = self.rngs
+        mem_shape = batch_dims + (self.hidden_features,)
+        h = self.carry_init(rngs.carry(), mem_shape, self.param_dtype)
+        return h
+
+    @property
+    def num_feature_axes(self) -> int:
+        return 1
+
+
+class RNN(Module):
+    """The ``RNN`` module takes any :class:`RNNCellBase` instance and applies it over a sequence
+
+    using :func:`flax.linen.scan`.
+    """
+
+    def __init__(
+        self,
+        cell: RNNCellBase,
+        time_major: bool = False,
+        return_carry: bool = False,
+        reverse: bool = False,
+        keep_order: bool = False,
+        unroll: int = 1,
+        rngs: rnglib.Rngs | None = None,
+    ):
+        self.cell = cell
+        self.time_major = time_major
+        self.return_carry = return_carry
+        self.reverse = reverse
+        self.keep_order = keep_order
+        self.unroll = unroll
+        if rngs is None:
+            rngs = rnglib.Rngs(0)
+        self.rngs = rngs
+
+    def __call__(
+        self,
+        inputs: Array,
+        *,
+        initial_carry: Carry | None = None,
+        seq_lengths: Array | None = None,
+        return_carry: bool | None = None,
+        time_major: bool | None = None,
+        reverse: bool | None = None,
+        keep_order: bool | None = None,
+        rngs: rnglib.Rngs | None = None,
+    ):
+        if return_carry is None:
+            return_carry = self.return_carry
+        if time_major is None:
+            time_major = self.time_major
+        if reverse is None:
+            reverse = self.reverse
+        if keep_order is None:
+            keep_order = self.keep_order
+
+        # Infer the number of batch dimensions from the input shape.
+        # Cells like ConvLSTM have additional spatial dimensions.
+        time_axis = 0 if time_major else inputs.ndim - (self.cell.num_feature_axes + 1)
+
+        # make time_axis positive
+        if time_axis < 0:
+            time_axis += inputs.ndim
+
+        if time_major:
+            # we add +1 because we moved the time axis to the front
+            batch_dims = inputs.shape[1 : -self.cell.num_feature_axes]
+        else:
+            batch_dims = inputs.shape[:time_axis]
+
+        # maybe reverse the sequence
+        if reverse:
+            inputs = jax.tree_util.tree_map(
+                lambda x: flip_sequences(
+                    x,
+                    seq_lengths,
+                    num_batch_dims=len(batch_dims),
+                    time_major=time_major,  # type: ignore
+                ),
+                inputs,
+            )
+        if rngs is None:
+            rngs = self.rngs
+        carry: Carry = (
+            self.cell.initialize_carry(
+                inputs.shape[:time_axis] + inputs.shape[time_axis + 1 :], rngs
+            )
+            if initial_carry is None
+            else initial_carry
+        )
+
+        slice_carry = seq_lengths is not None and return_carry
+
+        def scan_fn(cell: RNNCellBase, carry: Carry, x: Array) -> tuple[Carry, Array] | tuple[Carry, tuple[Carry, Array]]:
+            carry, y = cell(carry, x)
+            if slice_carry:
+                return carry, (carry, y)
+            return carry, y
+        state_axes = nnx.StateAxes({...: Carry}) # type: ignore[arg-type]
+        scan = nnx.scan(
+            scan_fn,
+            in_axes=(state_axes, Carry, time_axis),
+            out_axes=(Carry, (0, time_axis)) if slice_carry else (Carry, time_axis),
+            unroll=self.unroll,
+        )
+        scan_output = scan(self.cell, carry, inputs)
+
+        # Next we select the final carry. If a segmentation mask was provided and
+        # return_carry is True we slice the carry history and select the last valid
+        # carry for each sequence. Otherwise we just use the last carry.
+        if slice_carry:
+            assert seq_lengths is not None
+            _, (carries, outputs) = scan_output
+            # seq_lengths[None] expands the shape of the mask to match the
+            # number of dimensions of the carry.
+            carry = _select_last_carry(carries, seq_lengths)
+        else:
+            carry, outputs = scan_output
+
+        if reverse and keep_order:
+            outputs = jax.tree_util.tree_map(
+                lambda x: flip_sequences(
+                    x,
+                    seq_lengths,
+                    num_batch_dims=len(batch_dims),
+                    time_major=time_major,  # type: ignore
+                ),
+                outputs,
+            )
+
+        if return_carry:
+            return carry, outputs
+        else:
+            return outputs
+
+
+def _select_last_carry(sequence: A, seq_lengths: jnp.ndarray) -> A:
+    last_idx = seq_lengths - 1
+
+    def _slice_array(x: jnp.ndarray):
+        return x[last_idx, jnp.arange(x.shape[1])]
+
+    return jax.tree_util.tree_map(_slice_array, sequence)
+
+
+def _expand_dims_like(x, target):
+    """Expands the shape of `x` to match `target`'s shape by adding singleton dimensions."""
+    return x.reshape(list(x.shape) + [1] * (target.ndim - x.ndim))
+
+
+def flip_sequences(
+    inputs: Array,
+    seq_lengths: Array | None,
+    num_batch_dims: int,
+    time_major: bool,
+) -> Array:
+    """Flips a sequence of inputs along the time axis.
+
+    This function can be used to prepare inputs for the reverse direction of a
+    bidirectional LSTM. It solves the issue that, when naively flipping multiple
+    padded sequences stored in a matrix, the first elements would be padding
+    values for those sequences that were padded. This function keeps the padding
+    at the end, while flipping the rest of the elements.
+
+    Example:
+    ```python
+    inputs = [[1, 0, 0],
+              [2, 3, 0]
+              [4, 5, 6]]
+    lengths = [1, 2, 3]
+    flip_sequences(inputs, lengths) = [[1, 0, 0],
+                                       [3, 2, 0],
+                                       [6, 5, 4]]
+    ```
+
+    Args:
+      inputs: An array of input IDs <int>[batch_size, seq_length].
+      lengths: The length of each sequence <int>[batch_size].
+
+    Returns:
+      An ndarray with the flipped inputs.
+    """
+    # Compute the indices to put the inputs in flipped order as per above example.
+    time_axis = 0 if time_major else num_batch_dims
+    max_steps = inputs.shape[time_axis]
+
+    if seq_lengths is None:
+        # reverse inputs and return
+        inputs = jnp.flip(inputs, axis=time_axis)
+        return inputs
+
+    seq_lengths = jnp.expand_dims(seq_lengths, axis=time_axis)
+
+    # create indexes
+    idxs = jnp.arange(max_steps - 1, -1, -1)  # [max_steps]
+    if time_major:
+        idxs = jnp.reshape(idxs, [max_steps] + [1] * num_batch_dims)
+    else:
+        idxs = jnp.reshape(
+            idxs, [1] * num_batch_dims + [max_steps]
+        )  # [1, ..., max_steps]
+    idxs = (idxs + seq_lengths) % max_steps  # [*batch, max_steps]
+    idxs = _expand_dims_like(idxs, target=inputs)  # [*batch, max_steps, *features]
+    # Select the inputs in flipped order.
+    outputs = jnp.take_along_axis(inputs, idxs, axis=time_axis)
+
+    return outputs
+
+
+def _concatenate(a: Array, b: Array) -> Array:
+    """Concatenates two arrays along the last dimension."""
+    return jnp.concatenate([a, b], axis=-1)
+
+
+class RNNBase(Protocol):
+    def __call__(
+        self,
+        inputs: Array,
+        *,
+        initial_carry: Carry | None = None,
+        rngs: rnglib.Rngs | None = None,
+        seq_lengths: Array | None = None,
+        return_carry: bool | None = None,
+        time_major: bool | None = None,
+        reverse: bool | None = None,
+        keep_order: bool | None = None,
+    ) -> Output | tuple[Carry, Output]: ...
+
+
+class Bidirectional(Module):
+    """Processes the input in both directions and merges the results.
+
+    Example usage:
+
+    ```python
+    import nnx
+    import jax
+    import jax.numpy as jnp
+
+    # Define forward and backward RNNs
+    forward_rnn = RNN(GRUCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(0)))
+    backward_rnn = RNN(GRUCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(0)))
+
+    # Create Bidirectional layer
+    layer = Bidirectional(forward_rnn=forward_rnn, backward_rnn=backward_rnn)
+
+    # Input data
+    x = jnp.ones((2, 3, 3))
+
+    # Apply the layer
+    out = layer(x)
+    print(out.shape)
+    ```
+    """
+
+    forward_rnn: RNNBase
+    backward_rnn: RNNBase
+    merge_fn: Callable[[Array, Array], Array] = _concatenate
+    time_major: bool = False
+    return_carry: bool = False
+
+    def __init__(
+        self,
+        forward_rnn: RNNBase,
+        backward_rnn: RNNBase,
+        *,
+        merge_fn: Callable[[Array, Array], Array] = _concatenate,
+        time_major: bool = False,
+        return_carry: bool = False,
+        rngs: rnglib.Rngs | None = None,
+    ):
+        self.forward_rnn = forward_rnn
+        self.backward_rnn = backward_rnn
+        self.merge_fn = merge_fn
+        self.time_major = time_major
+        self.return_carry = return_carry
+        if rngs is None:
+            rngs = rnglib.Rngs(0)
+        self.rngs = rngs
+
+    def __call__(
+        self,
+        inputs: Array,
+        *,
+        initial_carry: tuple[Carry, Carry] | None = None,
+        rngs: rnglib.Rngs | None = None,
+        seq_lengths: Array | None = None,
+        return_carry: bool | None = None,
+        time_major: bool | None = None,
+        reverse: bool | None = None,  # unused
+        keep_order: bool | None = None,  # unused
+    ) -> Output | tuple[tuple[Carry, Carry], Output]:
+        if time_major is None:
+            time_major = self.time_major
+        if return_carry is None:
+            return_carry = self.return_carry
+        if rngs is None:
+            rngs = self.rngs
+        if initial_carry is not None:
+            initial_carry_forward, initial_carry_backward = initial_carry
+        else:
+            initial_carry_forward = None
+            initial_carry_backward = None
+        # Throw a warning in case the user accidentally re-uses the forward RNN
+        # for the backward pass and does not intend for them to share parameters.
+        if self.forward_rnn is self.backward_rnn:
+            logging.warning(
+                "forward_rnn and backward_rnn is the same object, so "
+                "they will share parameters."
+            )
+
+        # Encode in the forward direction.
+        carry_forward, outputs_forward = self.forward_rnn(
+            inputs,
+            initial_carry=initial_carry_forward,
+            rngs=rngs,
+            seq_lengths=seq_lengths,
+            return_carry=True,
+            time_major=time_major,
+            reverse=False,
+        )
+
+        # Encode in the backward direction.
+        carry_backward, outputs_backward = self.backward_rnn(
+            inputs,
+            initial_carry=initial_carry_backward,
+            rngs=rngs,
+            seq_lengths=seq_lengths,
+            return_carry=True,
+            time_major=time_major,
+            reverse=True,
+            keep_order=True,
+        )
+
+        carry = (carry_forward, carry_backward) if return_carry else None
+        outputs = jax.tree_util.tree_map(
+            self.merge_fn, outputs_forward, outputs_backward
+        )
+
+        if return_carry:
+            return carry, outputs
+        else:
+            return outputs
\ No newline at end of file
diff --git a/flax/nnx/statelib.py b/flax/nnx/statelib.py
index 4b3b1e387..59461fc46 100644
--- a/flax/nnx/statelib.py
+++ b/flax/nnx/statelib.py
@@ -21,6 +21,7 @@
 
 from flax.nnx import traversals
 from flax.nnx import filterlib, reprlib
+from flax.nnx import variablelib
 from flax.typing import PathParts
 
 A = tp.TypeVar('A')
@@ -424,4 +425,14 @@ def _split_state(
       # if we didn't break, set leaf to last state
       flat_states[-1][path] = value  # type: ignore[index] # mypy is wrong here?
 
-  return tuple(State.from_flat_path(flat_state) for flat_state in flat_states)
\ No newline at end of file
+  return tuple(State.from_flat_path(flat_state) for flat_state in flat_states)
+
+
+def create_path_filters(state: State):
+  flat_state = state.flat_state()
+  value_paths: dict[tp.Any, set[PathParts]] = {}
+  for path, value in flat_state.items():
+    if isinstance(value, (variablelib.Variable, variablelib.VariableState)):
+      value = value.value
+    value_paths.setdefault(value, set()).add(path)
+  return {filterlib.PathIn(*value_paths[value]): value for value in value_paths}
\ No newline at end of file
diff --git a/flax/nnx/transforms/compilation.py b/flax/nnx/transforms/compilation.py
index 88d99e8f7..e5ce20f8e 100644
--- a/flax/nnx/transforms/compilation.py
+++ b/flax/nnx/transforms/compilation.py
@@ -21,6 +21,8 @@
   extract,
   filterlib,
   graph,
+  statelib,
+  variablelib,
 )
 import jax
 import jax.core
@@ -36,13 +38,17 @@
 # -------------------------------
 
 
-class StateSharding:
+class StateSharding(extract.PrefixMapping):
   def __init__(
     self,
-    filter_sharding: tp.Mapping[filterlib.Filter, tp.Any]
+    filter_sharding: statelib.State
+    | tp.Mapping[filterlib.Filter, tp.Any]
     | tp.Iterable[tuple[filterlib.Filter, tp.Any]],
     /,
   ):
+    if isinstance(filter_sharding, statelib.State):
+      filter_sharding = statelib.create_path_filters(filter_sharding)  # type: ignore
+
     iterable = tuple(
       filter_sharding.items()
       if isinstance(filter_sharding, tp.Mapping)
@@ -59,6 +65,15 @@ def filters(self) -> tuple[filterlib.Filter, ...]:
   def shardings(self) -> tuple[tp.Any, ...]:
     return self._shardings
 
+  def map_prefix(
+    self, path: variablelib.PathParts, variable: variablelib.Variable
+  ) -> tp.Any:
+    for filter, sharding in zip(self.filters, self.shardings):
+      predicate = filterlib.to_predicate(filter)
+      if predicate(path, variable):
+        return sharding
+    raise ValueError(f'No axis found for {path=}, {variable=}')
+
   def __repr__(self):
     return f'StateSharding({dict(zip(self.filters, self.shardings))})'
 
diff --git a/flax/nnx/transforms/iteration.py b/flax/nnx/transforms/iteration.py
index c77e31951..3aaaa6ee3 100644
--- a/flax/nnx/transforms/iteration.py
+++ b/flax/nnx/transforms/iteration.py
@@ -20,7 +20,8 @@
 
 from flax import struct
 from flax.core.frozen_dict import FrozenDict
-from flax.nnx import extract, filterlib, graph, spmd
+from flax.nnx import extract, filterlib, graph, spmd, variablelib
+from flax.nnx import statelib
 from flax.nnx.module import Module
 from flax.nnx.statelib import State
 from flax.nnx.transforms.transforms import resolve_kwargs
@@ -54,16 +55,20 @@ class Carry:
 # -------------------------------
 
 
-class StateAxes:
+class StateAxes(extract.PrefixMapping):
 
   def __init__(
-      self,
-      filter_axes: (
-          tp.Mapping[filterlib.Filter, Index | type[Carry] | None]
-          | tp.Iterable[tuple[filterlib.Filter, Index | type[Carry] | None]]
-      ),
-      /,
+    self,
+    filter_axes: (
+      statelib.State
+      | tp.Mapping[filterlib.Filter, Index | type[Carry] | None]
+      | tp.Iterable[tuple[filterlib.Filter, Index | type[Carry] | None]]
+    ),
+    /,
   ):
+    if isinstance(filter_axes, statelib.State):
+      filter_axes = statelib.create_path_filters(filter_axes)  # type: ignore
+
     iterable = tuple(
         filter_axes.items()
         if isinstance(filter_axes, tp.Mapping)
@@ -80,6 +85,15 @@ def filters(self) -> tuple[filterlib.Filter, ...]:
   def axes(self) -> tuple[Index | type[Carry] | None, ...]:
     return self._axes
 
+  def map_prefix(
+    self, path: variablelib.PathParts, variable: variablelib.Variable
+  ) -> tp.Any:
+    for filter, axis in zip(self.filters, self.axes):
+      predicate = filterlib.to_predicate(filter)
+      if predicate(path, variable):
+        return axis
+    raise ValueError(f'No axis found for {path=}, {variable=}')
+
   def __repr__(self):
     return f'StateAxes({dict(self.items())})'
 
diff --git a/pyproject.toml b/pyproject.toml
index edcc4484f..8381caaab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -222,3 +222,7 @@ unfixable = []
 indent-style = "space"
 quote-style = "single"
 
+
+[tool.uv]
+# Ignore uv.lock and always upgrade the package to the latest
+upgrade-package = ["jax", "jaxlib", "orbax-checkpoint"]


Test patch:
diff --git a/tests/nnx/nn/recurrent_test.py b/tests/nnx/nn/recurrent_test.py
new file mode 100644
index 000000000..b724b69d7
--- /dev/null
+++ b/tests/nnx/nn/recurrent_test.py
@@ -0,0 +1,543 @@
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax, jax.numpy as jnp
+from jax import random
+
+from flax import linen
+from flax import nnx
+from flax.nnx.nn import initializers
+
+import numpy as np
+
+from absl.testing import absltest
+
+class TestLSTMCell(absltest.TestCase):
+    def test_basic(self):
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(0),
+        )
+        x = jnp.ones((2, 3))
+        carry = module.initialize_carry(x.shape, module.rngs)
+        new_carry, y = module(carry, x)
+        self.assertEqual(y.shape, (2, 4))
+
+    def test_lstm_sequence(self):
+        """Test LSTMCell over a sequence of inputs."""
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(0),
+        )
+        x = random.normal(random.PRNGKey(1), (5, 2, 3))  # seq_len, batch, feature
+        carry = module.initialize_carry(x.shape[1:], module.rngs)
+        outputs = []
+        for t in range(x.shape[0]):
+            carry, y = module(carry, x[t])
+            outputs.append(y)
+        outputs = jnp.stack(outputs)
+        self.assertEqual(outputs.shape, (5, 2, 4))
+
+    def test_lstm_with_different_dtypes(self):
+        """Test LSTMCell with different data types."""
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            dtype=jnp.bfloat16,
+            param_dtype=jnp.bfloat16,
+            rngs=nnx.Rngs(0),
+        )
+        x = jnp.ones((2, 3), dtype=jnp.bfloat16)
+        carry = module.initialize_carry(x.shape, module.rngs)
+        new_carry, y = module(carry, x)
+        self.assertEqual(y.dtype, jnp.bfloat16)
+        self.assertEqual(y.shape, (2, 4))
+
+    def test_lstm_with_custom_activations(self):
+        """Test LSTMCell with custom activation functions."""
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            gate_fn=jax.nn.relu,
+            activation_fn=jax.nn.elu,
+            rngs=nnx.Rngs(0),
+        )
+        x = jnp.ones((1, 3))
+        carry = module.initialize_carry(x.shape, module.rngs)
+        new_carry, y = module(carry, x)
+        self.assertEqual(y.shape, (1, 4))
+
+    def test_lstm_initialize_carry(self):
+        """Test the initialize_carry method."""
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            carry_init=initializers.ones,
+            rngs=nnx.Rngs(0),
+        )
+        x_shape = (1, 3)
+        carry = module.initialize_carry(x_shape, module.rngs)
+        c, h = carry
+        self.assertTrue(jnp.all(c == 1.0))
+        self.assertTrue(jnp.all(h == 1.0))
+        self.assertEqual(c.shape, (1, 4))
+        self.assertEqual(h.shape, (1, 4))
+
+    def test_lstm_with_variable_sequence_length(self):
+        """Test LSTMCell with variable sequence lengths."""
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(0)
+        )
+
+        # Simulate a batch with variable sequence lengths
+        x = jnp.array([
+            [[1, 2, 3], [4, 5, 6], [0, 0, 0]],  # Sequence length 2
+            [[7, 8, 9], [10, 11, 12], [13, 14, 15]],  # Sequence length 3
+        ])  # Shape: (batch_size=2, max_seq_length=3, features=3)
+
+        seq_lengths = jnp.array([2, 3])  # Actual lengths for each sequence
+        batch_size = x.shape[0]
+        max_seq_length = x.shape[1]
+        carry = module.initialize_carry((batch_size, 3), module.rngs)
+        outputs = []
+        for t in range(max_seq_length):
+            input_t = x[:, t, :]
+            carry, y = module(carry, input_t)
+            outputs.append(y)
+        outputs = jnp.stack(outputs, axis=1)  # Shape: (batch_size, max_seq_length, hidden_features)
+
+        # Zero out outputs beyond the actual sequence lengths
+        mask = (jnp.arange(max_seq_length)[None, :] < seq_lengths[:, None])
+        outputs = outputs * mask[:, :, None]
+        self.assertEqual(outputs.shape, (2, 3, 4))
+
+    def test_lstm_stateful(self):
+        """Test that LSTMCell maintains state across calls."""
+        module = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(0),
+        )
+        x1 = jnp.ones((1, 3))
+        x2 = jnp.ones((1, 3)) * 2
+        carry = module.initialize_carry(x1.shape)
+        carry, y1 = module(carry, x1)
+        carry, y2 = module(carry, x2)
+        self.assertEqual(y1.shape, (1, 4))
+        self.assertEqual(y2.shape, (1, 4))
+
+    def test_lstm_equivalence_with_flax_linen(self):
+        """Test that nnx.LSTMCell produces the same outputs as flax.linen.LSTMCell."""
+        in_features = 3
+        hidden_features = 4
+        key = random.PRNGKey(42)
+        x = random.normal(key, (1, in_features))
+
+        # Initialize nnx.LSTMCell
+        rngs_nnx = nnx.Rngs(0)
+        module_nnx = nnx.LSTMCell(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            rngs=rngs_nnx,
+        )
+        carry_nnx = module_nnx.initialize_carry(x.shape, rngs_nnx)
+        # Initialize flax.linen.LSTMCell
+        module_linen = linen.LSTMCell(
+            features=hidden_features,
+        )
+        carry_linen = module_linen.initialize_carry(random.PRNGKey(0), x.shape)
+        variables_linen = module_linen.init(random.PRNGKey(1), carry_linen, x)
+
+        # Copy parameters from flax.linen.LSTMCell to nnx.LSTMCell
+        params_linen = variables_linen['params']
+        # Map the parameters from linen to nnx
+        # Assuming the parameter names and shapes are compatible
+        # For a precise mapping, you might need to adjust parameter names
+        # Get the parameters from nnx module
+        nnx_params = module_nnx.__dict__
+
+        # Map parameters from linen to nnx
+        for gate in ['i', 'f', 'g', 'o']:
+            # Input kernels (input to gate)
+            if gate == 'f':
+                nnx_layer = getattr(module_nnx, f'if_')
+            else:
+                nnx_layer = getattr(module_nnx, f'i{gate}')
+            linen_params = params_linen[f'i{gate}']
+            nnx_layer.kernel.value = linen_params['kernel']
+            if nnx_layer.use_bias:
+                nnx_layer.bias.value = linen_params['bias']
+            # Hidden kernels (hidden state to gate)
+            nnx_layer = getattr(module_nnx, f'h{gate}')
+            linen_params = params_linen[f'h{gate}']
+            nnx_layer.kernel.value = linen_params['kernel']
+            if nnx_layer.use_bias:
+                nnx_layer.bias.value = linen_params['bias']
+
+        # Run both modules
+        new_carry_nnx, y_nnx = module_nnx(carry_nnx, x)
+        new_carry_linen, y_linen = module_linen.apply(variables_linen, carry_linen, x)
+
+        # Compare outputs
+        np.testing.assert_allclose(y_nnx, y_linen, atol=1e-5)
+        # Compare carries
+        for c_nnx, c_linen in zip(new_carry_nnx, new_carry_linen):
+            np.testing.assert_allclose(c_nnx, c_linen, atol=1e-5)
+
+class TestRNN(absltest.TestCase):
+
+    def test_rnn_with_lstm_cell(self):
+        """Test RNN module using LSTMCell."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(0),
+        )
+
+        # Initialize the RNN module with the LSTMCell
+        rnn = nnx.RNN(cell)
+
+        # Create input data (batch_size=2, seq_length=5, features=3)
+        x = jnp.ones((2, 5, 3))
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        self.assertEqual(outputs.shape, (2, 5, 4))  # Output features should match hidden_features
+
+    def test_rnn_with_gru_cell(self):
+        """Test RNN module using GRUCell."""
+        # Initialize the GRUCell
+        cell = nnx.GRUCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(1),
+        )
+
+        # Initialize the RNN module with the GRUCell
+        rnn = nnx.RNN(cell)
+
+        # Create input data (batch_size=2, seq_length=5, features=3)
+        x = jnp.ones((2, 5, 3))
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        self.assertEqual(outputs.shape, (2, 5, 4))  # Output features should match hidden_features
+
+    def test_rnn_time_major(self):
+        """Test RNN module with time_major=True."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(2),
+        )
+
+        # Initialize the RNN module with time_major=True
+        rnn = nnx.RNN(cell, time_major=True)
+
+        # Create input data (seq_length=5, batch_size=2, features=3)
+        x = jnp.ones((5, 2, 3))
+
+        # Initialize the carry
+        carry = cell.initialize_carry(x.shape[1:2] + x.shape[2:], cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        self.assertEqual(outputs.shape, (5, 2, 4))  # Output features should match hidden_features
+
+    def test_rnn_reverse(self):
+        """Test RNN module with reverse=True."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(3),
+        )
+
+        # Initialize the RNN module with reverse=True
+        rnn = nnx.RNN(cell, reverse=True)
+
+        # Create input data (batch_size=2, seq_length=5, features=3)
+        x = jnp.tile(jnp.arange(5), (2, 1)).reshape(2, 5, 1)  # Distinct values to check reversal
+        x = jnp.concatenate([x, x, x], axis=-1)  # Shape: (2, 5, 3)
+
+        # Run the RNN module
+        outputs = rnn(x)
+
+        # Check if the outputs are in reverse order
+        outputs_reversed = outputs[:, ::-1, :]
+        # Since we used distinct input values, we can compare outputs to check reversal
+        # For simplicity, just check the shapes here
+        self.assertEqual(outputs.shape, (2, 5, 4))
+        self.assertEqual(outputs_reversed.shape, (2, 5, 4))
+
+    def test_rnn_with_seq_lengths(self):
+        """Test RNN module with variable sequence lengths."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(4),
+        )
+
+        # Initialize the RNN module
+        rnn = nnx.RNN(cell, return_carry=True)
+
+        # Create input data with padding (batch_size=2, seq_length=5, features=3)
+        x = jnp.array([
+            [[1, 1, 1], [2, 2, 2], [3, 3, 3], [0, 0, 0], [0, 0, 0]],  # Sequence length 3
+            [[4, 4, 4], [5, 5, 5], [6, 6, 6], [7, 7, 7], [8, 8, 8]],  # Sequence length 5
+        ])  # Shape: (2, 5, 3)
+
+        seq_lengths = jnp.array([3, 5])  # Actual lengths for each sequence
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        final_carry, outputs = rnn(x, initial_carry=carry, seq_lengths=seq_lengths)
+
+        self.assertEqual(outputs.shape, (2, 5, 4))
+
+        self.assertEqual(final_carry[0].shape, (2, 4))  # c: (batch_size, hidden_features)
+        self.assertEqual(final_carry[1].shape, (2, 4))  # h: (batch_size, hidden_features)
+
+        # Todo: a better test by matching the outputs with the expected values
+
+    def test_rnn_with_keep_order(self):
+        """Test RNN module with reverse=True and keep_order=True."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(5),
+        )
+
+        # Initialize the RNN module with reverse=True and keep_order=True
+        rnn = nnx.RNN(cell, reverse=True, keep_order=True)
+
+        # Create input data (batch_size=2, seq_length=5, features=3)
+        x = jnp.tile(jnp.arange(5), (2, 1)).reshape(2, 5, 1)  # Distinct values to check reversal
+        x = jnp.concatenate([x, x, x], axis=-1)  # Shape: (2, 5, 3)
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        # Check if the outputs are in the original order despite processing in reverse
+        self.assertEqual(outputs.shape, (2, 5, 4))
+
+    def test_rnn_equivalence_with_flax_linen(self):
+        """Test that nnx.RNN produces the same outputs as flax.linen.RNN."""
+        in_features = 3
+        hidden_features = 4
+        seq_length = 5
+        batch_size = 2
+        key = random.PRNGKey(42)
+
+        # Create input data
+        x = random.normal(key, (batch_size, seq_length, in_features))
+
+        # Initialize nnx.LSTMCell and RNN
+        rngs_nnx = nnx.Rngs(0)
+        cell_nnx = nnx.LSTMCell(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            rngs=rngs_nnx,
+        )
+        rnn_nnx = nnx.RNN(cell_nnx)
+
+        # Initialize flax.linen.LSTMCell and RNN
+        cell_linen = linen.LSTMCell(features=hidden_features)
+        rnn_linen = linen.RNN(cell_linen)
+        carry_linen = cell_linen.initialize_carry(random.PRNGKey(0), x[:, 0].shape)
+        variables_linen = rnn_linen.init(random.PRNGKey(1), x)
+
+        # Copy parameters from flax.linen to nnx
+        params_linen = variables_linen['params']['cell']
+        # Copy cell parameters
+        for gate in ['i', 'f', 'g', 'o']:
+            # Input kernels
+            if gate == 'f':
+                nnx_layer = getattr(cell_nnx, f'if_')
+            else:
+                nnx_layer = getattr(cell_nnx, f'i{gate}')
+            linen_params = params_linen[f'i{gate}']
+            nnx_layer.kernel.value = linen_params['kernel']
+            if nnx_layer.use_bias:
+                nnx_layer.bias.value = linen_params['bias']
+            # Hidden kernels
+            nnx_layer = getattr(cell_nnx, f'h{gate}')
+            linen_params = params_linen[f'h{gate}']
+            nnx_layer.kernel.value = linen_params['kernel']
+            if nnx_layer.use_bias:
+                nnx_layer.bias.value = linen_params['bias']
+
+        # Initialize carries
+        carry_nnx = cell_nnx.initialize_carry((batch_size, in_features), rngs_nnx)
+
+        # Run nnx.RNN
+        outputs_nnx = rnn_nnx(x, initial_carry=carry_nnx)
+
+        # Run flax.linen.RNN
+        outputs_linen = rnn_linen.apply(variables_linen, x, initial_carry=carry_linen)
+
+        # Compare outputs
+        np.testing.assert_allclose(outputs_nnx, outputs_linen, atol=1e-5)
+
+    def test_rnn_with_unroll(self):
+        """Test RNN module with unroll parameter."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(6)
+        )
+
+        # Initialize the RNN module with unroll=2
+        rnn = nnx.RNN(cell, unroll=2)
+
+        # Create input data (batch_size=2, seq_length=6, features=3)
+        x = jnp.ones((2, 6, 3))
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        self.assertEqual(outputs.shape, (2, 6, 4))  # Output features should match hidden_features
+
+    def test_rnn_with_custom_cell(self):
+        """Test RNN module with a custom RNN cell."""
+        class CustomRNNCell(nnx.Module):
+            """A simple custom RNN cell."""
+
+            in_features: int
+            hidden_features: int
+
+            def __init__(self, in_features, hidden_features, rngs):
+                self.in_features = in_features
+                self.hidden_features = hidden_features
+                self.rngs = rngs
+                self.dense = nnx.Linear(
+                    in_features=in_features + hidden_features,
+                    out_features=hidden_features,
+                    rngs=rngs,
+                )
+
+            def __call__(self, carry, inputs):
+                h = carry
+                x = jnp.concatenate([inputs, h], axis=-1)
+                new_h = jax.nn.tanh(self.dense(x))
+                return new_h, new_h
+
+            def initialize_carry(self, input_shape, rngs):
+                batch_size = input_shape[0]
+                h = jnp.zeros((batch_size, self.hidden_features))
+                return h
+
+            @property
+            def num_feature_axes(self) -> int:
+                return 1
+
+        # Initialize the custom RNN cell
+        cell = CustomRNNCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(7)
+        )
+
+        # Initialize the RNN module
+        rnn = nnx.RNN(cell)
+
+        # Create input data (batch_size=2, seq_length=5, features=3)
+        x = jnp.ones((2, 5, 3))
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        self.assertEqual(outputs.shape, (2, 5, 4))  # Output features should match hidden_features
+
+    def test_rnn_with_different_dtypes(self):
+        """Test RNN module with different data types."""
+        # Initialize the LSTMCell with float16
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            dtype=jnp.float16,
+            param_dtype=jnp.float16,
+            rngs=nnx.Rngs(8),
+        )
+
+        # Initialize the RNN module
+        rnn = nnx.RNN(cell)
+
+        # Create input data (batch_size=2, seq_length=5, features=3)
+        x = jnp.ones((2, 5, 3), dtype=jnp.float16)
+
+        # Initialize the carry
+        carry = cell.initialize_carry((2, 3), cell.rngs)
+
+        # Run the RNN module
+        outputs = rnn(x, initial_carry=carry)
+
+        self.assertEqual(outputs.dtype, jnp.float16)
+        self.assertEqual(outputs.shape, (2, 5, 4))
+
+    def test_rnn_with_variable_batch_size(self):
+        """Test RNN module with variable batch sizes."""
+        # Initialize the LSTMCell
+        cell = nnx.LSTMCell(
+            in_features=3,
+            hidden_features=4,
+            rngs=nnx.Rngs(9),
+        )
+
+        # Initialize the RNN module
+        rnn = nnx.RNN(cell)
+
+        for batch_size in [1, 2, 5]:
+            # Create input data (batch_size, seq_length=5, features=3)
+            x = jnp.ones((batch_size, 5, 3))
+
+            # Initialize the carry
+            carry = cell.initialize_carry((batch_size, 3), cell.rngs)
+
+            # Run the RNN module
+            outputs = rnn(x, initial_carry=carry)
+
+            self.assertEqual(outputs.shape, (batch_size, 5, 4))
+
+if __name__ == '__main__':
+    absltest.main()
diff --git a/tests/nnx/transforms_test.py b/tests/nnx/transforms_test.py
index b29faeb48..5f478c432 100644
--- a/tests/nnx/transforms_test.py
+++ b/tests/nnx/transforms_test.py
@@ -2135,7 +2135,6 @@ def f(m):
 
   def test_consistent_aliasing_shared(self):
     class Shared(nnx.Module):
-
       def __init__(self):
         self.a = nnx.Param(jnp.zeros((3, 3)))
 
@@ -2148,17 +2147,46 @@ def __init__(self, shared: Shared):
     m1 = Foo(shared)
     m2 = Foo(shared)
 
-    @partial(nnx.vmap, in_axes=(0, 1))
+    @nnx.vmap(in_axes=(0, 1))
     def f(m1, m2):
       pass
 
     with self.assertRaisesRegex(
-        ValueError,
-        r'Inconsistent aliasing detected([\s\S]*)Shared([\s\S]*)a:'
-        r' 0([\s\S]*)a: 1',
+      ValueError,
+      r'Inconsistent aliasing detected([\s\S]*)Param([\s\S]*)a:'
+      r' 0([\s\S]*)a: 1',
     ):
       f(m1, m2)
 
+  def test_equivalent_state_axes_mapping(self):
+    m = nnx.Linear(3, 3, rngs=nnx.Rngs(0))
+
+    sa1 = nnx.StateAxes({...: 0})
+    sa2 = nnx.StateAxes({nnx.Param: 0})
+
+    @nnx.vmap(in_axes=(0, sa1, sa2))
+    def f(m1, m2, m3):
+      pass
+
+    f(m, m, m)
+
+  def test_equivalent_state_sharding_mapping(self):
+    m = nnx.Linear(3, 3, rngs=nnx.Rngs(0))
+
+    mesh = jax.sharding.Mesh(jax.devices(), ('mp',))
+    sharding = jax.sharding.NamedSharding(
+      mesh, jax.sharding.PartitionSpec('mp')
+    )
+
+    sa1 = nnx.StateSharding({...: sharding})
+    sa2 = nnx.StateSharding({nnx.Param: sharding})
+
+    @nnx.jit(in_shardings=(sharding, sa1, sa2))
+    def f(m1, m2, m3):
+      pass
+
+    f(m, m, m)
+
   @absltest.skip('Enable once jax#19586 resolved')
   def test_captured_module_in_return_error(self):
     class Foo(nnx.Module):
@@ -2317,6 +2345,44 @@ def create_block(rngs: nnx.Rngs):
     self.assertEqual(m.kernel.value.shape, (5, 16, 32))
     self.assertEqual(m.kernel.sharding, ('c', 'a', 'b'))
 
+  def test_state_axes_from_state(self):
+    class Model(nnx.Module):
+      def __init__(self, din, dout, *, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(din, dout, rngs=rngs)
+        self.bn = nnx.BatchNorm(dout, rngs=rngs)
+
+    model = Model(2, 3, rngs=nnx.Rngs(0))
+    state = nnx.state(model)
+
+    state['linear']['kernel'] = 0
+    state['linear']['bias'] = 1
+    state['bn']['scale'] = 0
+    state['bn']['mean'] = 1
+    state['bn']['var'] = 0
+    state['bn']['bias'] = None
+
+    state_axes = nnx.StateAxes(state)
+
+    self.assertEqual(state_axes.map_prefix(('linear', 'kernel'), None), 0)
+    self.assertEqual(state_axes.map_prefix(('linear', 'bias'), None), 1)
+    self.assertEqual(state_axes.map_prefix(('bn', 'scale'), None), 0)
+    self.assertEqual(state_axes.map_prefix(('bn', 'mean'), None), 1)
+    self.assertEqual(state_axes.map_prefix(('bn', 'var'), None), 0)
+    self.assertEqual(state_axes.map_prefix(('bn', 'bias'), None), None)
+
+    @nnx.vmap(out_axes=state_axes, axis_size=5)
+    def create_block():
+      return Model(2, 3, rngs=nnx.Rngs(0))
+
+    model = create_block()
+
+    self.assertEqual(model.linear.kernel.shape, (5, 2, 3))
+    self.assertEqual(model.linear.bias.shape, (3, 5))
+    self.assertEqual(model.bn.scale.shape, (5, 3))
+    self.assertEqual(model.bn.mean.shape, (3, 5))
+    self.assertEqual(model.bn.var.shape, (5, 3))
+    self.assertEqual(model.bn.bias.shape, (3,))
+
 
 class TestPmap(absltest.TestCase):