Merge pull request #433 from nasyxx:patch-1

PiperOrigin-RevId: 478994245
google-deepmind · Oct 5, 2022 · e026a15 · e026a15
2 parents 232c48d + 71e7455
commit e026a15
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/docs/api.rst b/docs/api.rst
@@ -10,6 +10,8 @@ Common Optimizers
     adagrad
     adam
     adamw
+    adamax
+    adamaxw
     fromage
     lamb
     lars
@@ -43,6 +45,16 @@ Adam
 
 .. autofunction:: adam
 
+Adamax
+~~~~
+
+.. autofunction:: adamax
+
+AdamaxW
+~~~~~
+
+.. autofunction:: adamaxw
+
 AdamW
 ~~~~~
 
@@ -147,6 +159,7 @@ Gradient Transforms
     Params
     scale
     scale_by_adam
+    scale_by_adamax
     scale_by_belief
     scale_by_factored_rms
     scale_by_optimistic_gradient
@@ -257,6 +270,7 @@ Optax Transforms and States
 
 .. autofunction:: scale
 .. autofunction:: scale_by_adam
+.. autofunction:: scale_by_adamax
 .. autofunction:: scale_by_belief
 .. autofunction:: scale_by_factored_rms
 .. autofunction:: scale_by_param_block_norm

diff --git a/optax/__init__.py b/optax/__init__.py
@@ -19,6 +19,8 @@
 from optax._src.alias import adafactor
 from optax._src.alias import adagrad
 from optax._src.alias import adam
+from optax._src.alias import adamax
+from optax._src.alias import adamaxw
 from optax._src.alias import adamw
 from optax._src.alias import dpsgd
 from optax._src.alias import fromage
@@ -122,6 +124,7 @@
 from optax._src.transform import EmaState
 from optax._src.transform import scale
 from optax._src.transform import scale_by_adam
+from optax._src.transform import scale_by_adamax
 from optax._src.transform import scale_by_belief
 from optax._src.transform import scale_by_optimistic_gradient
 from optax._src.transform import scale_by_param_block_norm
@@ -175,6 +178,8 @@
     "adafactor",
     "adagrad",
     "adam",
+    "adamax",
+    "adamaxw",
     "adamw",
     "adaptive_grad_clip",
     "AdaptiveGradClipState",
@@ -267,6 +272,7 @@
     "safe_root_mean_squares",
     "ScalarOrSchedule",
     "scale_by_adam",
+    "scale_by_adamax",
     "scale_by_belief",
     "scale_by_factored_rms",
     "scale_by_param_block_norm",