Let momo accept additional extra arguments.

PiperOrigin-RevId: 627842385
google-deepmind · Apr 24, 2024 · fbd8f5d · fbd8f5d
1 parent d9c2bbe
commit fbd8f5d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/docs/api/contrib.rst b/docs/api/contrib.rst
@@ -16,6 +16,10 @@ Experimental features and algorithms that don't meet the
     dpsgd
     mechanize
     MechanicState
+    momo
+    MomoState
+    momo_adam
+    MomoAdamState
     prodigy
     ProdigyState
     sam
@@ -63,6 +67,13 @@ Mechanize
 .. autoclass:: MechanicState
     :members:
 
+Momo
+~~~~
+.. autofunction:: momo
+.. autoclass:: MomoState
+.. autofunction:: momo_adam
+.. autoclass:: MomoAdamState
+
 Prodigy
 ~~~~~~~
 .. autofunction:: prodigy

diff --git a/optax/contrib/_momo.py b/optax/contrib/_momo.py
@@ -96,8 +96,11 @@ def update_fn(
       updates: base.Updates,
       state: MomoState,
       params: Optional[base.Params],
+      *,
       value: Optional[Array] = None,
+      **extra_args,
   ) -> tuple[base.Updates, MomoState]:
+    del extra_args
     if params is None:
       raise ValueError(base.NO_PARAMS_MSG)
     if value is None:
@@ -227,8 +230,11 @@ def update_fn(
       updates: base.Updates,
       state: MomoAdamState,
       params: Optional[base.Params],
+      *,
       value: Optional[Array],
+      **extra_args,
   ) -> tuple[base.Updates, MomoAdamState]:
+    del extra_args
     if params is None:
       raise ValueError(base.NO_PARAMS_MSG)
     if value is None: