Repository: google/flax

Issue:
Dev container is outdated
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- OS Platform and Distribution: MacOS
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: Not applicable, problems occurs when building the dev containers.
- Python version: 3.12
- GPU/TPU model and memory: Not Applicable
- CUDA version (if applicable): Not Applicable
- Docker desktop version: 4.28.0
- Docker engine version: 25.0.3
- VS code version: 1.94.2

### Problem you have encountered:
Unable to build the devcontainer.

### What you expected to happen:
I expected to connect to the dev container after building it, but the instruction at ./dev/README.md aren't up to date. Specifically  the configuration files. 

### Logs, error messages, etc:
1. `ERROR: Ignored the following versions that require a different python version ERROR: Could not find a version that satisfies the requirement jaxlib  (from versions: none)`
2. `The 'sklearn' PyPI package is deprecated, use 'scikit-learn'`
3. `file:///workspaces/flax/flax does not appear to be a Python project: neither 'setup.py' nor 'pyproject.toml' found.`


### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.
1. Copy ./dev/.devcontainer to the root directory of the repo.
2. Open root directory in VS code. 
3. Launch the container when prompted by VS Code.
4. Check that it fails to run.

Hints from issue comments:

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/flax_gspmd.ipynb b/docs_nnx/guides/flax_gspmd.ipynb
index 4011c56c2..5787ab750 100644
--- a/docs_nnx/guides/flax_gspmd.ipynb
+++ b/docs_nnx/guides/flax_gspmd.ipynb
@@ -46,7 +46,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -56,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -74,14 +74,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 +106,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [
     {
@@ -140,7 +140,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -202,7 +202,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
@@ -249,7 +249,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
@@ -296,7 +296,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
     {
@@ -309,31 +309,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 +349,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": {},
@@ -419,37 +413,37 @@
   },
   {
    "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 +452,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 +498,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)"
    ]
@@ -530,7 +518,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
@@ -543,33 +531,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": {},
@@ -598,7 +584,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
@@ -647,14 +633,14 @@
   },
   {
    "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"
      ]
     }
    ],
@@ -683,7 +669,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 27,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -728,37 +714,37 @@
   },
   {
    "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 +753,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": {},
@@ -870,7 +850,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..424520f24 100644
--- a/docs_nnx/guides/flax_gspmd.md
+++ b/docs_nnx/guides/flax_gspmd.md
@@ -232,7 +232,7 @@ 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)
 ```
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/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/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):