Fix issue with flat_state()

Author: vfdev-5

docs/source/JAX_Vision_transformer.ipynb

@@ -295,7 +295,7 @@
     "\n",
     "    # Notice the use of `flax.nnx.state`.\n",
     "    flax_model_params = nnx.state(dst_model, nnx.Param)\n",
-    "    flax_model_params_fstate = flax_model_params.flat_state()\n",
+    "    flax_model_params_fstate = dict(flax_model_params.flat_state())\n",
     "\n",
     "    # Mapping from Flax parameter names to TF parameter names.\n",
     "    params_name_mapping = {\n",

docs/source/JAX_Vision_transformer.md

@@ -251,7 +251,7 @@ def vit_inplace_copy_weights(*, src_model, dst_model):
 
     # Notice the use of `flax.nnx.state`.
     flax_model_params = nnx.state(dst_model, nnx.Param)
-    flax_model_params_fstate = flax_model_params.flat_state()
+    flax_model_params_fstate = dict(flax_model_params.flat_state())
 
     # Mapping from Flax parameter names to TF parameter names.
     params_name_mapping = {

docs/source/JAX_image_captioning.ipynb

@@ -687,7 +687,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "id": "a15b8ac4-6cf1-4d4a-a888-840230ea93d1",
    "metadata": {},
    "outputs": [],
@@ -700,7 +700,7 @@
     "    tf_model_params_fstate = nnx.traversals.flatten_mapping(tf_model_params)\n",
     "\n",
     "    flax_model_params = nnx.state(dst_model, nnx.Param)\n",
-    "    flax_model_params_fstate = flax_model_params.flat_state()\n",
+    "    flax_model_params_fstate = dict(flax_model_params.flat_state())\n",
     "\n",
     "    src_num_params = sum([p.size for p in tf_model_params_fstate.values()])\n",
     "    dst_num_params = sum([p.value.size for p in flax_model_params_fstate.values()])\n",
@@ -1281,12 +1281,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "id": "f8b6dba9-41e1-4d03-9296-ad25af1132fa",
    "metadata": {},
    "outputs": [],
    "source": [
-    "for key in list(nnx.state(model, trainable_params_filter).flat_state().keys()):\n",
+    "for key, _ in nnx.state(model, trainable_params_filter).flat_state():\n",
     "    assert \"encoder\" not in key"
    ]
   },

docs/source/JAX_image_captioning.md

@@ -458,7 +458,7 @@ def vit_inplace_copy_weights(*, src_model, dst_model):
     tf_model_params_fstate = nnx.traversals.flatten_mapping(tf_model_params)
 
     flax_model_params = nnx.state(dst_model, nnx.Param)
-    flax_model_params_fstate = flax_model_params.flat_state()
+    flax_model_params_fstate = dict(flax_model_params.flat_state())
 
     src_num_params = sum([p.size for p in tf_model_params_fstate.values()])
     dst_num_params = sum([p.value.size for p in flax_model_params_fstate.values()])
@@ -913,7 +913,7 @@ model_diffstate = nnx.DiffState(0, trainable_params_filter)
 ```
 
 ```{code-cell} ipython3
-for key in list(nnx.state(model, trainable_params_filter).flat_state().keys()):
+for key, _ in nnx.state(model, trainable_params_filter).flat_state():
     assert "encoder" not in key
 ```
 

docs/source/JAX_porting_PyTorch_model.ipynb

@@ -180,13 +180,13 @@ class Model(nnx.Module):
 To inspect NNX module's learnable parameters and buffers, we can use `nnx.state`:
 ```python
 nnx_module = ...
-for k, v in nnx.state(nnx_module, nnx.Param).flat_state().items():
+for k, v in nnx.state(nnx_module, nnx.Param).flat_state():
     print(
         k,
         v.value.mean() if v.value is not None else None
     )
 
-for k, v in nnx.state(nnx_module, (nnx.BatchStat, Buffer)).flat_state().items():
+for k, v in nnx.state(nnx_module, (nnx.BatchStat, Buffer)).flat_state():
     print(
         k,
         v.value.mean() if v.value.dtype == "float32" else v.value.sum()