[nnx] add Using Filters guide

google · Jun 28, 2024 · 0aa5fe6 · 0aa5fe6
1 parent f610e6c
commit 0aa5fe6
Show file tree

Hide file tree

Showing 7 changed files with 702 additions and 5 deletions.
diff --git a/docs/api_reference/flax.nnx/filterlib.rst b/docs/api_reference/flax.nnx/filterlib.rst
@@ -0,0 +1,16 @@
+filterlib
+------------------------
+
+.. automodule:: flax.nnx
+.. currentmodule:: flax.nnx
+
+
+.. autofunction:: flax.nnx.filterlib.to_predicate
+.. autoclass:: WithTag
+.. autoclass:: PathContains
+.. autoclass:: OfType
+.. autoclass:: Any
+.. autoclass:: All
+.. autoclass:: Not
+.. autoclass:: Everything
+.. autoclass:: Nothing
diff --git a/docs/api_reference/flax.nnx/index.rst b/docs/api_reference/flax.nnx/index.rst
@@ -17,4 +17,5 @@ Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/nnx/i
   variables
   helpers
   visualization
+  filterlib