update: test cases

Author: kozistr

tests/constants.py

@@ -30,6 +30,7 @@
     AdaBound,
     AdaDelta,
     AdaFactor,
+    AdaGC,
     AdaHessian,
     Adai,
     Adalite,
@@ -86,6 +87,7 @@
     SGDSaI,
     Shampoo,
     SignSGD,
+    SimplifiedAdEMAMix,
     SophiaH,
     StableAdamW,
     StableSPAM,
@@ -513,6 +515,7 @@
     (AdamG, {'lr': 1e0, 'p': 0.5}, 60),
     (AdEMAMix, {'lr': 1e0}, 3),
     (AdEMAMix, {'lr': 1e0, 't_alpha_beta3': 5}, 3),
+    (SimplifiedAdEMAMix, {'lr': 1e0, 'beta1_warmup': 2, 'weight_decay': 1e-3}, 3),
     (
         SOAP,
         {'lr': 1e0, 'shampoo_beta': 0.95, 'precondition_frequency': 1, 'merge_dims': False, 'precondition_1d': True},
@@ -567,6 +570,7 @@
     (EXAdam, {'lr': 1e-1, 'weight_decay': 1e-3}, 5),
     (SCION, {'lr': 5e-1, 'constraint': False, 'weight_decay': 1e-3}, 5),
     (SCION, {'lr': 1e-1, 'constraint': True, 'lmo_type': 'col_norm'}, 10),
+    (AdaGC, {'lr': 1e-1, 'warmup_steps': 2, 'weight_decay': 1e-3}, 5),
     (Ranger25, {'lr': 1e-1}, 3),
     (Ranger25, {'lr': 1e-1, 't_alpha_beta3': 5}, 3),
     (Ranger25, {'lr': 5e-2, 'stable_adamw': False, 'orthograd': False, 'eps': None, 'lookahead_merge_time': 2}, 3),

tests/test_load_modules.py

@@ -34,7 +34,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 97
+    assert len(get_supported_optimizers()) == 99
     assert len(get_supported_optimizers('adam*')) == 8
     assert len(get_supported_optimizers(['adam*', 'ranger*'])) == 11