Merge pull request #4715 from google:nnx-mutable-array-p1

PiperOrigin-RevId: 756050354

Author: Flax Authors

docs_nnx/conf.py

@@ -139,7 +139,7 @@
 
 # -- Options for myst ----------------------------------------------
 # uncomment line below to avoid running notebooks during development
-nb_execution_mode = 'off'
+# nb_execution_mode = 'off'
 # Notebook cell execution timeout; defaults to 30.
 nb_execution_timeout = 100
 # List of patterns, relative to source directory, that match notebook
@@ -151,6 +151,7 @@
   'flax/nnx',  # exclude nnx
   'guides/demo.ipynb',  # TODO(cgarciae): broken, remove or update
   'guides/gemma.ipynb',
+  'guides/bridge_guide.ipynb',  # TODO(cgarciae): broken, bridge doesn't support Linen sow yet
 ]
 # raise exceptions on execution so CI can catch errors
 nb_execution_allow_errors = False

docs_nnx/guides/bridge_guide.ipynb

@@ -284,7 +284,7 @@
     "  predicates = [nnx.filterlib.to_predicate(f) for f in filters]\n",
     "  flat_states: list[dict[KeyPath, Any]] = [{} for p in predicates]\n",
     "\n",
-    "  for path, value in state.flat_state():\n",
+    "  for path, value in state:\n",
     "    for i, predicate in enumerate(predicates):\n",
     "      if predicate(path, value):\n",
     "        flat_states[i][path] = value\n",
@@ -415,7 +415,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

docs_nnx/guides/bridge_guide.md

@@ -149,7 +149,7 @@ def split(node, *filters):
   predicates = [nnx.filterlib.to_predicate(f) for f in filters]
   flat_states: list[dict[KeyPath, Any]] = [{} for p in predicates]
 
-  for path, value in state.flat_state():
+  for path, value in state:
     for i, predicate in enumerate(predicates):
       if predicate(path, value):
         flat_states[i][path] = value

docs_nnx/guides/filters_guide.ipynb

@@ -147,7 +147,7 @@
     "\n",
     "for _ in range(10):\n",
     "  x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\n",
-    "  state, loss = jax_train_step(graphdef, state, x, y)\n",
+    "  loss, state = jax_train_step(graphdef, state, x, y)\n",
     "\n",
     "# update objects after training\n",
     "nnx.update((model, optimizer, metrics), state)"

docs_nnx/guides/filters_guide.md

@@ -119,7 +119,7 @@ def jax_train_step(graphdef, state, x, y):
 
 for _ in range(10):
   x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))
-  state, loss = jax_train_step(graphdef, state, x, y)
+  loss, state = jax_train_step(graphdef, state, x, y)
 
 # update objects after training
 nnx.update((model, optimizer, metrics), state)