update: Kate optimizer

Author: kozistr

pytorch_optimizer/__init__.py

@@ -66,6 +66,7 @@
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.gravity import Gravity
 from pytorch_optimizer.optimizer.grokfast import GrokFastAdamW, gradfilter_ema, gradfilter_ma
+from pytorch_optimizer.optimizer.kate import Kate
 from pytorch_optimizer.optimizer.lamb import Lamb
 from pytorch_optimizer.optimizer.lars import LARS
 from pytorch_optimizer.optimizer.lion import Lion
@@ -199,6 +200,7 @@
     ScheduleFreeAdamW,
     FAdam,
     GrokFastAdamW,
+    Kate,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

tests/constants.py

@@ -50,6 +50,7 @@
     GaLore,
     Gravity,
     GrokFastAdamW,
+    Kate,
     Lamb,
     Lion,
     Nero,
@@ -461,6 +462,7 @@
     (ScheduleFreeAdamW, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (FAdam, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (GrokFastAdamW, {'lr': 1e0, 'weight_decay': 1e-3}, 10),
+    (Kate, {'lr': 5e-2}, 10),
 ]
 ADANORM_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adanorm': True}, 10),

tests/test_load_modules.py

@@ -38,7 +38,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 68
+    assert len(get_supported_optimizers()) == 69
 
 
 def test_get_supported_lr_schedulers():