have TrainState subclass nnx.Module

google · Apr 1, 2024 · 543686b · 543686b
1 parent f2bdcd8
commit 543686b
Show file tree

Hide file tree

Showing 12 changed files with 666 additions and 197 deletions.
diff --git a/docs/api_reference/flax.experimental.nnx/index.rst b/docs/api_reference/flax.experimental.nnx/index.rst
@@ -10,6 +10,7 @@ Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/exper
   nn/index
   rnglib
   spmd
+  training/index
   transforms
   variables
 
diff --git a/docs/api_reference/flax.experimental.nnx/training/index.rst b/docs/api_reference/flax.experimental.nnx/training/index.rst
@@ -0,0 +1,11 @@
+flax.experimental.nnx.nnx.training
+----------------------------
+
+Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/experimental/nnx/index.html>`__ for more details.
+
+.. toctree::
+  :maxdepth: 3
+
+  metrics
+  optimizer
+
diff --git a/docs/api_reference/flax.experimental.nnx/training/metrics.rst b/docs/api_reference/flax.experimental.nnx/training/metrics.rst
@@ -0,0 +1,15 @@
+Metrics
+------------------------
+
+.. automodule:: flax.experimental.nnx.metrics
+.. currentmodule:: flax.experimental.nnx.metrics
+
+.. autoclass:: Metric
+   :members:
+.. autoclass:: Average
+   :members:
+.. autoclass:: Accuracy
+   :members:
+.. autoclass:: MultiMetric
+   :members:
+
diff --git a/docs/api_reference/flax.experimental.nnx/training/optimizer.rst b/docs/api_reference/flax.experimental.nnx/training/optimizer.rst
@@ -0,0 +1,8 @@
+Optimizer
+------------------------
+
+.. automodule:: flax.experimental.nnx.optimizer
+.. currentmodule:: flax.experimental.nnx.optimizer
+
+.. autoclass:: Optimizer
+   :members: