pytorch-optimizer 3.1.1__tar.gz → 3.1.2__tar.gz

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: pytorch_optimizer
-Version: 3.1.1
+Version: 3.1.2
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
 Home-page: https://github.com/kozistr/pytorch_optimizer
 License: Apache-2.0
-Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,FAdam,Fromage,GaLore,Gravity,GrokFast,GSAM,Kate,Lamb,LARS,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,StableAdamW,SWATS,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,FAdam,Fromage,GaLore,Gravity,GrokFast,GSAM,Kate,Lamb,LARS,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,StableAdamW,SWATS,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -49,7 +49,7 @@ Description-Content-Type: text/markdown
 
 **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch. 
 I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.  
-Currently, **75 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **76 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -212,6 +212,7 @@ supported_optimizers = get_supported_optimizers()
 | AdamMini      | *Use Fewer Learning Rates To Gain More*                                                           | [github](https://github.com/zyushun/Adam-mini)                                                                 | <https://arxiv.org/abs/2406.16793>                                                         | [cite](https://github.com/zyushun/Adam-mini?tab=readme-ov-file#citation)                                                            |
 | TRAC          | *Adaptive Parameter-free Optimization*                                                            | [github](https://github.com/ComputationalRobotics/TRAC)                                                        | <https://arxiv.org/abs/2405.16642>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240516642M/exportcitation)                                                        |
 | AdamG         | *Towards Stability of Parameter-free Optimization*                                                |                                                                                                                | <https://arxiv.org/abs/2405.04376>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240504376P/exportcitation)                                                        |
+| AdEMAMix      | *Better, Faster, Older*                                                                           | [github](https://github.com/nanowell/AdEMAMix-Optimizer-Pytorch)                                               | <https://arxiv.org/abs/2409.03137>                                                         | [cite](https://github.com/nanowell/AdEMAMix-Optimizer-Pytorch?tab=readme-ov-file#reference)                                         |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/README.md

@@ -10,7 +10,7 @@
 
 **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch. 
 I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.  
-Currently, **75 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **76 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -173,6 +173,7 @@ supported_optimizers = get_supported_optimizers()
 | AdamMini      | *Use Fewer Learning Rates To Gain More*                                                           | [github](https://github.com/zyushun/Adam-mini)                                                                 | <https://arxiv.org/abs/2406.16793>                                                         | [cite](https://github.com/zyushun/Adam-mini?tab=readme-ov-file#citation)                                                            |
 | TRAC          | *Adaptive Parameter-free Optimization*                                                            | [github](https://github.com/ComputationalRobotics/TRAC)                                                        | <https://arxiv.org/abs/2405.16642>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240516642M/exportcitation)                                                        |
 | AdamG         | *Towards Stability of Parameter-free Optimization*                                                |                                                                                                                | <https://arxiv.org/abs/2405.04376>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240504376P/exportcitation)                                                        |
+| AdEMAMix      | *Better, Faster, Older*                                                                           | [github](https://github.com/nanowell/AdEMAMix-Optimizer-Pytorch)                                               | <https://arxiv.org/abs/2409.03137>                                                         | [cite](https://github.com/nanowell/AdEMAMix-Optimizer-Pytorch?tab=readme-ov-file#reference)                                         |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.1.1"
+version = "3.1.2"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -11,15 +11,15 @@ repository = "https://github.com/kozistr/pytorch_optimizer"
 documentation = "https://pytorch-optimizers.readthedocs.io/en/latest"
 keywords = [
     "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound",
-    "AdaDelta", "AdaFactor", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai",
-    "Adalite", "AdaLomo", "AdamMini", "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "AvaGrad",
-    "bSAM", "CAME", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DiffGrad", "FAdam",
-    "Fromage", "GaLore", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD",
-    "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
-    "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "SGDP", "Shampoo", "ScalableShampoo", "SGDW",
-    "SignSGD", "SM3", "SopihaH", "SRMM", "StableAdamW", "SWATS", "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal",
-    "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky",
-    "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "AdaDelta", "AdaFactor", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdEMAMix", "AdaHessian",
+    "Adai", "Adalite", "AdaLomo", "AdamMini", "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo",
+    "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DiffGrad",
+    "FAdam", "Fromage", "GaLore", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LARS", "Lion", "LOMO", "Lookahead",
+    "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam",
+    "Ranger", "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "SGDP", "Shampoo",
+    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "StableAdamW", "SWATS", "Tiger", "TRAC", "WSAM",
+    "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky",
+    "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
@@ -71,7 +71,7 @@ lint.select = [
 ]
 lint.ignore = [
     "B905", "D100", "D102", "D104", "D105", "D107", "D203", "D213", "D413", "PIE790", "PLR0912", "PLR0913", "PLR0915",
-    "PLR2004", "RUF013", "Q003",
+    "PLR2004", "RUF013", "Q003", "ARG002",
 ]
 lint.fixable = ["ALL"]
 lint.unfixable = ["F401"]

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/__init__.py

@@ -54,6 +54,7 @@ from pytorch_optimizer.optimizer.adanorm import AdaNorm
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
 from pytorch_optimizer.optimizer.adashift import AdaShift
 from pytorch_optimizer.optimizer.adasmooth import AdaSmooth
+from pytorch_optimizer.optimizer.ademamix import AdEMAMix
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.aggmo import AggMo
 from pytorch_optimizer.optimizer.aida import Aida
@@ -208,6 +209,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     AdamMini,
     AdaLOMO,
     AdamG,
+    AdEMAMix,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/lr_scheduler/rex.py

@@ -52,7 +52,7 @@ class REXScheduler(LRScheduler):
 
         return self.min_lr + (self.max_lr - self.min_lr) * ((1.0 - progress) / (1.0 - progress / 2.0))
 
-    def step(self, epoch: Optional[int] = None) -> float:  # noqa: ARG002
+    def step(self, epoch: Optional[int] = None) -> float:
         value: float = self.get_linear_lr()
 
         self.step_t += 1

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/a2grad.py

@@ -27,6 +27,7 @@ class A2Grad(BaseOptimizer):
         lips: float = 10.0,
         rho: float = 0.5,
         variant: str = 'uni',
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(lips, 'lips')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adabelief.py

@@ -42,6 +42,7 @@ class AdaBelief(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-16,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adabound.py

@@ -37,6 +37,7 @@ class AdaBound(BaseOptimizer):
         ams_bound: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adadelta.py

@@ -26,6 +26,7 @@ class AdaDelta(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         eps: float = 1e-6,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(rho, 'rho', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adafactor.py

@@ -49,6 +49,7 @@ class AdaFactor(BaseOptimizer):
         eps1: float = 1e-30,
         eps2: float = 1e-3,
         momentum_dtype: torch.dtype = torch.bfloat16,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adahessian.py

@@ -40,6 +40,7 @@ class AdaHessian(BaseOptimizer):
         hessian_distribution: HUTCHINSON_G = 'rademacher',
         adam_debias: bool = False,
         eps: float = 1e-16,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adai.py

@@ -35,6 +35,7 @@ class Adai(BaseOptimizer):
         dampening: float = 1.0,
         use_gc: bool = False,
         eps: float = 1e-3,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adalite.py

@@ -35,6 +35,7 @@ class Adalite(BaseOptimizer):
         tau: float = 1.0,
         eps1: float = 1e-6,
         eps2: float = 1e-10,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adam_mini.py

@@ -38,6 +38,7 @@ class AdamMini(BaseOptimizer):  # pragma: no cover
         num_heads: int = 32,
         num_query_groups: Optional[int] = None,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adamax.py

@@ -32,6 +32,7 @@ class AdaMax(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adamg.py

@@ -32,6 +32,7 @@ class AdamG(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adamod.py

@@ -31,6 +31,7 @@ class AdaMod(BaseOptimizer):
         fixed_decay: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adamp.py

@@ -45,6 +45,7 @@ class AdamP(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adams.py

@@ -36,6 +36,7 @@ class AdamS(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adamw.py

@@ -29,6 +29,7 @@ class StableAdamW(BaseOptimizer):
         weight_decay: float = 1e-2,
         weight_decouple: bool = True,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adan.py

@@ -37,6 +37,7 @@ class Adan(BaseOptimizer):
         r: float = 0.95,
         adanorm: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adanorm.py

@@ -34,6 +34,7 @@ class AdaNorm(BaseOptimizer):
         ams_bound: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adapnm.py

@@ -36,6 +36,7 @@ class AdaPNM(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adashift.py

@@ -28,6 +28,7 @@ class AdaShift(BaseOptimizer):
         keep_num: int = 10,
         reduce_func: Optional[Callable] = torch.max,
         eps: float = 1e-10,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/adasmooth.py

@@ -26,6 +26,7 @@ class AdaSmooth(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         eps: float = 1e-6,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

pytorch_optimizer-3.1.2/pytorch_optimizer/optimizer/ademamix.py

@@ -0,0 +1,149 @@
+import math
+from typing import Optional
+
+import torch
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class AdEMAMix(BaseOptimizer):
+    r"""Better, Faster, Older.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param fixed_decay: bool. fix weight decay.
+    :param alpha: float. usually between 4 and 10 would work well.
+    :param t_alpha_beta3: Optional[float]. total number of iterations is preferred when needed.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.9, 0.999, 0.9999),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = False,
+        fixed_decay: bool = False,
+        alpha: float = 5.0,
+        t_alpha_beta3: Optional[float] = None,
+        eps: float = 1e-8,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(alpha, 'alpha')
+        self.validate_non_negative(t_alpha_beta3, 't_alpha_beta3')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'alpha': alpha,
+            't_alpha_beta3': t_alpha_beta3,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'AdEMAMix'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = torch.zeros_like(p)
+                state['exp_avg_slow'] = torch.zeros_like(p)
+
+    @staticmethod
+    def schedule_alpha(t_alpha_beta3: Optional[float], step: int, alpha: float) -> float:
+        if t_alpha_beta3 is None:
+            return alpha
+        return min(step * alpha / t_alpha_beta3, alpha)
+
+    @staticmethod
+    def schedule_beta3(t_alpha_beta3: Optional[float], step: int, beta1: float, beta3: float) -> float:
+        if t_alpha_beta3 is None:
+            return beta3
+
+        log_beta1, log_beta3 = math.log(beta1), math.log(beta3)
+
+        return min(
+            math.exp(
+                log_beta1 * log_beta3 / ((1.0 - step / t_alpha_beta3) * log_beta3 + (step / t_alpha_beta3) * log_beta1)
+            ),
+            beta3,
+        )
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2, beta3 = group['betas']
+
+            bias_correction1: float = self.debias(beta1, group['step'])
+            bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+
+            alpha_t: float = self.schedule_alpha(group['t_alpha_beta3'], group['step'], group['alpha'])
+            beta3_t: float = self.schedule_beta3(group['t_alpha_beta3'], group['step'], beta1, beta3)
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['exp_avg_slow'] = torch.zeros_like(p)
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                exp_avg, exp_avg_sq, exp_avg_slow = state['exp_avg'], state['exp_avg_sq'], state['exp_avg_slow']
+
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+                exp_avg_slow.mul_(beta3_t).add_(grad, alpha=1.0 - beta3_t)
+
+                de_nom = (exp_avg_sq.sqrt() / bias_correction2_sq).add_(group['eps'])
+
+                step_size = group['lr'] / bias_correction1
+
+                p.addcdiv_(exp_avg + alpha_t * exp_avg_slow, de_nom, value=-step_size)
+
+        return loss

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/aggmo.py

@@ -24,6 +24,7 @@ class AggMo(BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         fixed_decay: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/aida.py

@@ -46,6 +46,7 @@ class Aida(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/alig.py

@@ -25,6 +25,7 @@ class AliG(BaseOptimizer):
         projection_fn: Optional[Callable] = None,
         momentum: float = 0.0,
         adjusted_momentum: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(max_lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/amos.py

@@ -31,6 +31,7 @@ class Amos(BaseOptimizer):
         c_coef: float = 0.25,
         d_coef: float = 0.25,
         eps: float = 1e-18,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0, range_type='[)')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/apollo.py

@@ -33,6 +33,7 @@ class Apollo(BaseOptimizer):
         weight_decay_type: str = 'l2',
         warmup_steps: int = 500,
         eps: float = 1e-4,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[]')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/avagrad.py

@@ -30,6 +30,7 @@ class AvaGrad(BaseOptimizer):
         fixed_decay: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-1,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/came.py

@@ -35,6 +35,7 @@ class CAME(BaseOptimizer):
         ams_bound: bool = False,
         eps1: float = 1e-30,
         eps2: float = 1e-16,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/dadapt.py

@@ -39,6 +39,7 @@ class DAdaptAdaGrad(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         eps: float = 0.0,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0, range_type='[)')
@@ -266,6 +267,7 @@ class DAdaptAdam(BaseOptimizer):
         fixed_decay: bool = False,
         bias_correction: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
@@ -423,6 +425,7 @@ class DAdaptSGD(BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         fixed_decay: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0, range_type='[)')
@@ -560,6 +563,7 @@ class DAdaptAdan(BaseOptimizer):
         d0: float = 1e-6,
         growth_rate: float = float('inf'),
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
@@ -721,6 +725,7 @@ class DAdaptLion(BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         fixed_decay: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/diffgrad.py

@@ -40,6 +40,7 @@ class DiffGrad(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/fadam.py

@@ -30,6 +30,7 @@ class FAdam(BaseOptimizer):
         eps: float = 1e-8,
         momentum_dtype: torch.dtype = torch.float32,
         fim_dtype: torch.dtype = torch.float32,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/fromage.py

@@ -22,7 +22,7 @@ class Fromage(BaseOptimizer):
         norms to lie within 2x their initial norms. This regularises the model class.
     """
 
-    def __init__(self, params: PARAMETERS, lr: float = 1e-2, p_bound: Optional[float] = None):
+    def __init__(self, params: PARAMETERS, lr: float = 1e-2, p_bound: Optional[float] = None, **kwargs):
         self.validate_learning_rate(lr)
 
         self.p_bound = p_bound

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/galore.py

@@ -21,7 +21,12 @@ class GaLoreProjector:
     """
 
     def __init__(
-        self, rank: int = 128, update_proj_gap: int = 50, scale: float = 1.0, projection_type: PROJECTION_TYPE = 'std'
+        self,
+        rank: int = 128,
+        update_proj_gap: int = 50,
+        scale: float = 1.0,
+        projection_type: PROJECTION_TYPE = 'std',
+        **kwargs,
     ):
         self.rank = rank
         self.update_proj_gap = update_proj_gap

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/gravity.py

@@ -20,6 +20,7 @@ class Gravity(BaseOptimizer):
         lr: float = 1e-2,
         alpha: float = 0.01,
         beta: float = 0.9,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(alpha, 'alpha', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/grokfast.py

@@ -128,6 +128,7 @@ class GrokFastAdamW(BaseOptimizer):
         fixed_decay: bool = False,
         normalize_lr: bool = True,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/kate.py

@@ -26,6 +26,7 @@ class Kate(BaseOptimizer):
         weight_decouple: bool = True,
         fixed_decay: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(delta, 'delta', 0.0, 1.0, '[)')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/lamb.py

@@ -53,6 +53,7 @@ class Lamb(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-6,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/lars.py

@@ -26,6 +26,7 @@ class LARS(BaseOptimizer):
         dampening: float = 0.0,
         trust_coefficient: float = 1e-3,
         nesterov: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(weight_decay, 'weight_decay')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/lion.py

@@ -31,6 +31,7 @@ class Lion(BaseOptimizer):
         use_gc: bool = False,
         r: float = 0.95,
         adanorm: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/lomo.py

@@ -30,6 +30,7 @@ class LOMO(BaseOptimizer):
         lr: float = 1e-3,
         clip_grad_norm: Optional[float] = None,
         clip_grad_value: Optional[float] = None,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(clip_grad_norm, 'clip_grad_norm')
@@ -228,6 +229,7 @@ class AdaLOMO(BaseOptimizer):
         clip_grad_value: Optional[float] = None,
         eps1: float = 1e-30,
         eps2: float = 1e-3,
+        **kwargs,
     ) -> None:  # fmt: skip
         self.validate_learning_rate(lr)
         self.validate_non_negative(weight_decay, 'weight_decay')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/lookahead.py

@@ -22,6 +22,7 @@ class Lookahead(BaseOptimizer):
         k: int = 5,
         alpha: float = 0.5,
         pullback_momentum: str = 'none',
+        **kwargs,
     ) -> None:
         self.validate_positive(k, 'k')
         self.validate_range(alpha, 'alpha', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/madgrad.py

@@ -32,6 +32,7 @@ class MADGRAD(BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         eps: float = 1e-6,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(weight_decay, 'weight_decay')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/msvag.py

@@ -13,7 +13,13 @@ class MSVAG(BaseOptimizer):
     :param beta: float. Moving average (momentum) constant (scalar tensor or float value).
     """
 
-    def __init__(self, params: PARAMETERS, lr: float = 1e-2, beta: float = 0.9):
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-2,
+        beta: float = 0.9,
+        **kwargs,
+    ):
         self.validate_learning_rate(lr)
         self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[]')
 

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/nero.py

@@ -17,7 +17,13 @@ class Nero(BaseOptimizer):
     """
 
     def __init__(
-        self, params: PARAMETERS, lr: float = 0.01, beta: float = 0.999, constraints: bool = True, eps: float = 1e-8
+        self,
+        params: PARAMETERS,
+        lr: float = 0.01,
+        beta: float = 0.999,
+        constraints: bool = True,
+        eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[]')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/novograd.py

@@ -32,6 +32,7 @@ class NovoGrad(BaseOptimizer):
         grad_averaging: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/padam.py

@@ -30,6 +30,7 @@ class PAdam(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/pid.py

@@ -30,6 +30,7 @@ class PID(BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         fixed_decay: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/pnm.py

@@ -28,6 +28,7 @@ class PNM(BaseOptimizer):
         weight_decouple: bool = True,
         fixed_decay: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/prodigy.py

@@ -44,6 +44,7 @@ class Prodigy(BaseOptimizer):
         bias_correction: bool = False,
         safeguard_warmup: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas((*betas, beta3))

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/qhadam.py

@@ -30,6 +30,7 @@ class QHAdam(BaseOptimizer):
         weight_decouple: bool = False,
         fixed_decay: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/qhm.py

@@ -26,6 +26,7 @@ class QHM(BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         fixed_decay: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/radam.py

@@ -36,6 +36,7 @@ class RAdam(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/ranger.py

@@ -43,6 +43,7 @@ class Ranger(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-5,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/ranger21.py

@@ -82,6 +82,7 @@ class Ranger21(BaseOptimizer):
         norm_loss_factor: float = 1e-4,
         adam_debias: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_learning_rate(warm_down_min_lr)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/rotograd.py

@@ -176,7 +176,7 @@ class RotateOnly(nn.Module):
         backbone: nn.Module,
         heads: Sequence[nn.Module],
         latent_size: int,
-        *args,  # noqa: ARG002
+        *args,
         burn_in_period: int = 20,
         normalize_losses: bool = False,
     ):

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/schedulefree.py

@@ -36,6 +36,7 @@ class ScheduleFreeSGD(BaseOptimizer):
         weight_lr_power: float = 2.0,
         warmup_steps: int = 0,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0, range_type='[]')
@@ -180,6 +181,7 @@ class ScheduleFreeAdamW(BaseOptimizer):
         warmup_steps: int = 0,
         ams_bound: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/sgd.py

@@ -27,6 +27,7 @@ class AccSGD(BaseOptimizer):
         xi: float = 10.0,
         constant: float = 0.7,
         weight_decay: float = 0.0,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(kappa, 'kappa')
@@ -124,6 +125,7 @@ class SGDW(BaseOptimizer):
         weight_decouple: bool = True,
         dampening: float = 0.0,
         nesterov: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0)
@@ -222,6 +224,7 @@ class ASGD(BaseOptimizer):
         theta: float = 1.0,
         dampening: float = 1.0,
         eps: float = 1e-5,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(amplifier, 'amplifier')
@@ -339,6 +342,7 @@ class SignSGD(BaseOptimizer):
         momentum: float = 0.9,
         weight_decay: float = 0.0,
         weight_decouple: bool = True,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'beta', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/sgdp.py

@@ -36,6 +36,7 @@ class SGDP(BaseOptimizer):
         wd_ratio: float = 0.1,
         nesterov: bool = False,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_non_negative(weight_decay, 'weight_decay')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/shampoo.py

@@ -36,6 +36,7 @@ class Shampoo(BaseOptimizer):
         fixed_decay: bool = False,
         preconditioning_compute_steps: int = 1,
         matrix_eps: float = 1e-6,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0)
@@ -211,6 +212,7 @@ class ScalableShampoo(BaseOptimizer):
         diagonal_eps: float = 1e-10,
         matrix_eps: float = 1e-6,
         use_svd: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/shampoo_utils.py

@@ -37,7 +37,7 @@ class Graft:
         r"""Get preconditioned gradient."""
         return grad
 
-    def update_momentum(self, update: torch.Tensor, unused_beta1: float) -> torch.Tensor:  # noqa: ARG002
+    def update_momentum(self, update: torch.Tensor, unused_beta1: float) -> torch.Tensor:
         r"""Update momentum."""
         return update
 

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/sm3.py

@@ -23,6 +23,7 @@ class SM3(BaseOptimizer):
         momentum: float = 0.0,
         beta: float = 0.0,
         eps: float = 1e-30,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(momentum, 'momentum', 0.0, 1.0)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/sophia.py

@@ -38,6 +38,7 @@ class SophiaH(BaseOptimizer):
         num_samples: int = 1,
         hessian_distribution: HUTCHINSON_G = 'gaussian',
         eps: float = 1e-12,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/srmm.py

@@ -16,7 +16,14 @@ class SRMM(BaseOptimizer):
     :param memory_length: Optional[int]. internal memory length for moving average. None for no refreshing.
     """
 
-    def __init__(self, params: PARAMETERS, lr: float = 0.01, beta: float = 0.5, memory_length: Optional[int] = 100):
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 0.01,
+        beta: float = 0.5,
+        memory_length: Optional[int] = 100,
+        **kwargs,
+    ):
         self.validate_learning_rate(lr)
         self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[]')
 

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/swats.py

@@ -38,6 +38,7 @@ class SWATS(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-6,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/tiger.py

@@ -24,6 +24,7 @@ class Tiger(BaseOptimizer):
         weight_decay: float = 0.01,
         weight_decouple: bool = True,
         fixed_decay: bool = False,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[)')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/trac.py

@@ -105,6 +105,7 @@ class TRAC(BaseOptimizer):
         num_coefs: int = 128,
         s_prev: float = 1e-8,
         eps: float = 1e-8,
+        **kwargs,
     ):
         self.validate_positive(num_coefs, 'num_coefs')
         self.validate_non_negative(s_prev, 's_prev')

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.1.2}/pytorch_optimizer/optimizer/yogi.py

@@ -36,6 +36,7 @@ class Yogi(BaseOptimizer):
         adanorm: bool = False,
         adam_debias: bool = False,
         eps: float = 1e-3,
+        **kwargs,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)