added to common tests

optax/contrib/_common_test.py

@@ -34,6 +34,8 @@
     dict(opt_name='cocob', opt_kwargs=dict(alpha=100.0, eps=1e-8)),
     dict(opt_name='dadapt_adamw', opt_kwargs=dict(learning_rate=1e-1)),
     dict(opt_name='prodigy', opt_kwargs=dict(learning_rate=1e-1)),
+    dict(opt_name='momo', opt_kwargs=dict(learning_rate=1e-1)),
+    dict(opt_name='momo_adam', opt_kwargs=dict(learning_rate=1e-1)),
 )
 
 
@@ -42,7 +44,7 @@ def _setup_parabola(dtype):
   initial_params = jnp.array([-1.0, 10.0, 1.0], dtype=dtype)
   final_params = jnp.array([1.0, -1.0, 1.0], dtype=dtype)
 
-  @jax.grad
+  @jax.value_and_grad
   def get_updates(params):
     return jnp.sum(numerics.abs_sq(params - final_params))
 
@@ -57,7 +59,7 @@ def _setup_rosenbrock(dtype):
   initial_params = jnp.array([0.0, 0.0], dtype=dtype)
   final_params = jnp.array([a, a**2], dtype=dtype)
 
-  @jax.grad
+  @jax.value_and_grad
   def get_updates(params):
     return numerics.abs_sq(a - params[0]) + b * numerics.abs_sq(
         params[1] - params[0] ** 2
@@ -79,8 +81,12 @@ def test_optimizers(self, opt_name, opt_kwargs, target, dtype):
 
     @jax.jit
     def step(params, state):
-      updates = get_updates(params)
-      updates, state = opt.update(updates, state, params)
+      loss, updates = get_updates(params)
+      if opt_name in ['momo', 'momo_adam']:
+        update_kwargs = {'loss': loss}
+      else:
+        update_kwargs = {}
+      updates, state = opt.update(updates, state, params, **update_kwargs)
       params = update.apply_updates(params, updates)
       return params, state
 
@@ -107,12 +113,20 @@ def test_optimizers_can_be_wrapped_in_inject_hyperparams(
     params = [jnp.negative(jnp.ones((2, 3))), jnp.ones((2, 5, 2))]
     grads = [jnp.ones((2, 3)), jnp.negative(jnp.ones((2, 5, 2)))]
 
+    if opt_name in ['momo', 'momo_adam']:
+      update_kwargs = {'loss': jnp.array(0.)}
+    else:
+      update_kwargs = {}
+
     state = self.variant(opt.init)(params)
-    updates, new_state = self.variant(opt.update)(grads, state, params)
+    updates, new_state = self.variant(opt.update)(
+      grads, state, params, **update_kwargs
+    )
 
     state_inject = self.variant(opt_inject.init)(params)
     updates_inject, new_state_inject = self.variant(opt_inject.update)(
-        grads, state_inject, params)
+        grads, state_inject, params, **update_kwargs
+    )
 
     with self.subTest('Equality of updates.'):
       chex.assert_trees_all_close(updates_inject, updates, rtol=1e-4)

optax/contrib/momo.py

@@ -39,7 +39,7 @@ def momo(
     learning_rate: base.ScalarOrSchedule = 1.0,
     beta: float = 0.9,
     lb: float = 0.0,
-    weight_decay: float = 0.
+    weight_decay: float = 0.0
 ) -> base.GradientTransformationExtraArgs:
   """Adaptive Learning Rates for SGD with momentum.
 
@@ -134,11 +134,11 @@ class MomoAdamState(NamedTuple):
 
 
 def momo_adam(
-    learning_rate: base.ScalarOrSchedule = 1.0,
+    learning_rate: base.ScalarOrSchedule = 1e-2,
     betas: tuple[float, float] = (0.9, 0.999),
     eps: float = 1e-8,
     lb: float = 0.0,
-    weight_decay: float = 0.
+    weight_decay: float = 0.0
 ) -> base.GradientTransformationExtraArgs:
   """Adaptive Learning Rates for Adam(W).
 
@@ -175,7 +175,7 @@ def init_fn(params: base.Params) -> MomoAdamState:
     exp_avg_sq = tu.tree_map(lambda p: jnp.zeros(p.shape, jnp.float32), params)
     barf = 0
     gamma = 0
-    count = 0
+    count = jnp.zeros([], jnp.int32)
     return MomoAdamState(exp_avg, exp_avg_sq, barf, gamma, count)
 
   def update_fn(

optax/contrib/momo_test.py

@@ -91,7 +91,7 @@ class MomoAdamTest(chex.TestCase):
       dtype=(jnp.float32,),
   )
   def test_optimization(self, opt_name, target, dtype):
-    opt = getattr(contrib, opt_name)()
+    opt = getattr(contrib, opt_name)(learning_rate=0.1)
     initial_params, final_params, get_updates = target(dtype)
     @jax.jit
     def step(params, state):