Detail that learning_rate can be scalar or a schedule.

The doc was slightly misleading see #399. PiperOrigin-RevId: 604372850
google-deepmind · Feb 6, 2024 · d2bb5f1 · d2bb5f1
1 parent 7907220
commit d2bb5f1
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 237 deletions.
diff --git a/docs/api/transformations.rst b/docs/api/transformations.rst
@@ -43,6 +43,7 @@ Transformations
     ScaleByBeliefState
     scale_by_factored_rms
     FactoredState
+    scale_by_learning_rate
     scale_by_lion
     ScaleByLionState
     scale_by_novograd