Merge pull request #4897 from google:pytree-guide

PiperOrigin-RevId: 803235639

Author: Flax Authors

.github/workflows/flax_test.yml

@@ -82,7 +82,7 @@ jobs:
 
   tests:
     name: Run Tests
-    needs: [pre-commit, commit-count, test-import]
+    needs: [pre-commit, commit-count]
     runs-on: ubuntu-24.04-16core
     strategy:
       # Make sure to change `github_check_runs` in `copy.bara.sky` if you change the tests here.

docs_nnx/api_reference/flax.nnx/helpers.rst

@@ -7,5 +7,9 @@ helpers
 
 .. autoclass:: Sequential
    :members:
+.. autoclass:: List
+   :members:
+.. autoclass:: Dict
+   :members:
 .. autoclass:: TrainState
    :members:

docs_nnx/api_reference/flax.nnx/object.rst

@@ -4,10 +4,16 @@ object
 .. automodule:: flax.nnx
 .. currentmodule:: flax.nnx
 
+.. autoclass:: Pytree
+   :members:
 .. autoclass:: Object
    :members:
 .. autofunction:: data
 .. autodata:: Data
    :annotation:
-.. autofunction:: is_data_type
-.. autofunction:: register_data_type
+.. autofunction:: static
+.. autodata:: Static
+   :annotation:
+.. autofunction:: is_data
+.. autofunction:: register_data_type
+.. autofunction:: check_pytree

docs_nnx/api_reference/flax.nnx/rnglib.rst

@@ -9,4 +9,5 @@ rnglib
 .. autoclass:: RngStream
    :members:
 .. autofunction:: split_rngs
+.. autofunction:: fork_rngs
 .. autofunction:: reseed

docs_nnx/guides/array_ref.ipynb

@@ -150,6 +150,14 @@
     "print(f\"{variable.has_ref = }\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "839332be",
+   "metadata": {},
+   "source": [
+    "Mention `nnx.use_refs` can be used as global flag"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "1b2632f1",

docs_nnx/guides/array_ref.md

@@ -58,6 +58,10 @@ with nnx.use_refs(True):
 print(f"{variable.has_ref = }")
 ```
 
+Mention `nnx.use_refs` can be used as global flag
+
++++
+
 ### Changing Status
 
 ```{code-cell} ipython3