pytorch-optimizer 2.9.0__tar.gz → 2.10.0__tar.gz

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: pytorch-optimizer
-Version: 2.9.0
-Summary: optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas.
+Version: 2.10.0
+Summary: optimizer & lr scheduler 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,AdaMod,AdaNorm,AdaPNM,AdaSmooth,Adai,AdamP,AdamS,Adan,AggMo,AliG,Apollo,AvaGrad,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DiffGrad,Fromage,Gravity,LARS,Lamb,Lion,MADGRAD,MSVAG,Nero,NovoGrad,PID,PNM,QHAdam,QHM,RAdam,Ranger,Ranger21,SGDP,SGDW,SM3,SRMM,SWATS,ScalableShampoo,Shampoo,Yogi,SAM,GSAM,PCGrad,RotoGrad
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,AdamP,AdamS,Adan,AggMo,AliG,Amos,Apollo,AvaGrad,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DiffGrad,Fromage,Gravity,LARS,Lamb,Lion,MADGRAD,MSVAG,Nero,NovoGrad,PID,PNM,QHAdam,QHM,RAdam,Ranger,Ranger21,SGDP,SGDW,SignSGD,SM3,SopihaH,SRMM,SWATS,ScalableShampoo,Shampoo,Yogi,SAM,GSAM,PCGrad,RotoGrad
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -21,13 +21,8 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.6
 Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
@@ -59,7 +54,7 @@ pytorch-optimizer
 
 | **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, 50 optimizers, 6 lr schedulers are supported!
+| Currently, 54 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -172,9 +167,9 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | Adai         | *Disentangling the Effects of Adaptive Learning Rate and Momentum*                                | `github <https://github.com/zeke-xie/adaptive-inertia-adai>`__                    | `https://arxiv.org/abs/2006.15815 <https://arxiv.org/abs/2006.15815>`__                       | `cite <https://github.com/zeke-xie/adaptive-inertia-adai#citing>`__                                                  |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                                        |
+| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                                        |
+| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | GSAM         | *Surrogate Gap Guided Sharpness-Aware Minimization*                                               | `github <https://github.com/juntang-zhuang/GSAM>`__                               | `https://openreview.net/pdf?id=edONMAnhLu- <https://openreview.net/pdf?id=edONMAnhLu->`__     | `cite <https://github.com/juntang-zhuang/GSAM#citation>`__                                                           |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
@@ -256,6 +251,14 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | AdaDelta     | *An Adaptive Learning Rate Method*                                                                |                                                                                   | `https://arxiv.org/abs/1212.5701v1 <https://arxiv.org/abs/1212.5701v1>`__                     | `cite <https://ui.adsabs.harvard.edu/abs/2012arXiv1212.5701Z/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Amos         | *An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale*                 | `github <https://github.com/google-research/jestimator>`__                        | `https://arxiv.org/abs/2210.11693 <https://arxiv.org/abs/2210.11693>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2022arXiv221011693T/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| SignSGD      | *Compressed Optimisation for Non-Convex Problems*                                                 | `github <https://github.com/jxbz/signSGD>`__                                      | `https://arxiv.org/abs/1802.04434 <https://arxiv.org/abs/1802.04434>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2018arXiv180204434B/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| AdaHessian   | *An Adaptive Second Order Optimizer for Machine Learning*                                         | `github <https://github.com/amirgholami/adahessian>`__                            | `https://arxiv.org/abs/2006.00719 <https://arxiv.org/abs/2006.00719>`__                       | `cite <https://github.com/amirgholami/adahessian#citation>`__                                                        |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Sophia       | *A Scalable Stochastic Second-order Optimizer for Language Model Pre-training*                    | `github <https://github.com/Liuhong99/Sophia>`__                                  | `https://arxiv.org/abs/2305.14342 <https://arxiv.org/abs/2305.14342>`__                       | `cite <https://github.com/Liuhong99/Sophia>`__                                                                       |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/README.rst

@@ -16,7 +16,7 @@ pytorch-optimizer
 
 | **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, 50 optimizers, 6 lr schedulers are supported!
+| Currently, 54 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -129,9 +129,9 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | Adai         | *Disentangling the Effects of Adaptive Learning Rate and Momentum*                                | `github <https://github.com/zeke-xie/adaptive-inertia-adai>`__                    | `https://arxiv.org/abs/2006.15815 <https://arxiv.org/abs/2006.15815>`__                       | `cite <https://github.com/zeke-xie/adaptive-inertia-adai#citing>`__                                                  |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                                        |
+| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                                        |
+| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | GSAM         | *Surrogate Gap Guided Sharpness-Aware Minimization*                                               | `github <https://github.com/juntang-zhuang/GSAM>`__                               | `https://openreview.net/pdf?id=edONMAnhLu- <https://openreview.net/pdf?id=edONMAnhLu->`__     | `cite <https://github.com/juntang-zhuang/GSAM#citation>`__                                                           |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
@@ -213,6 +213,14 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | AdaDelta     | *An Adaptive Learning Rate Method*                                                                |                                                                                   | `https://arxiv.org/abs/1212.5701v1 <https://arxiv.org/abs/1212.5701v1>`__                     | `cite <https://ui.adsabs.harvard.edu/abs/2012arXiv1212.5701Z/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Amos         | *An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale*                 | `github <https://github.com/google-research/jestimator>`__                        | `https://arxiv.org/abs/2210.11693 <https://arxiv.org/abs/2210.11693>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2022arXiv221011693T/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| SignSGD      | *Compressed Optimisation for Non-Convex Problems*                                                 | `github <https://github.com/jxbz/signSGD>`__                                      | `https://arxiv.org/abs/1802.04434 <https://arxiv.org/abs/1802.04434>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2018arXiv180204434B/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| AdaHessian   | *An Adaptive Second Order Optimizer for Machine Learning*                                         | `github <https://github.com/amirgholami/adahessian>`__                            | `https://arxiv.org/abs/2006.00719 <https://arxiv.org/abs/2006.00719>`__                       | `cite <https://github.com/amirgholami/adahessian#citation>`__                                                        |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| Sophia       | *A Scalable Stochastic Second-order Optimizer for Language Model Pre-training*                    | `github <https://github.com/Liuhong99/Sophia>`__                                  | `https://arxiv.org/abs/2305.14342 <https://arxiv.org/abs/2305.14342>`__                       | `cite <https://github.com/Liuhong99/Sophia>`__                                                                       |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.9.0"
-description = "optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
+version = "2.10.0"
+description = "optimizer & lr scheduler collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
 maintainers = ["kozistr <kozistr@gmail.com>"]
@@ -9,7 +9,7 @@ readme = "README.rst"
 homepage = "https://github.com/kozistr/pytorch_optimizer"
 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", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "Adai", "AdamP", "AdamS", "Adan", "AggMo", "AliG", "Apollo", "AvaGrad", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "LARS", "Lamb", "Lion", "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PID", "PNM", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "SGDP", "SGDW", "SM3", "SRMM", "SWATS", "ScalableShampoo", "Shampoo", "Yogi", "SAM", "GSAM", "PCGrad", "RotoGrad"]
+keywords = ["pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound", "AdaDelta", "AdaFactor", "AdaMax", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai", "AdamP", "AdamS", "Adan", "AggMo", "AliG", "Amos", "Apollo", "AvaGrad", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "LARS", "Lamb", "Lion", "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PID", "PNM", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "SGDP", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "SWATS", "ScalableShampoo", "Shampoo", "Yogi", "SAM", "GSAM", "PCGrad", "RotoGrad"]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Development Status :: 5 - Production/Stable",

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/__init__.py

@@ -21,6 +21,7 @@ from pytorch_optimizer.optimizer.adabelief import AdaBelief
 from pytorch_optimizer.optimizer.adabound import AdaBound
 from pytorch_optimizer.optimizer.adadelta import AdaDelta
 from pytorch_optimizer.optimizer.adafactor import AdaFactor
+from pytorch_optimizer.optimizer.adahessian import AdaHessian
 from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adamax import AdaMax
 from pytorch_optimizer.optimizer.adamod import AdaMod
@@ -34,6 +35,7 @@ from pytorch_optimizer.optimizer.adasmooth import AdaSmooth
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.aggmo import AggMo
 from pytorch_optimizer.optimizer.alig import AliG
+from pytorch_optimizer.optimizer.amos import Amos
 from pytorch_optimizer.optimizer.apollo import Apollo
 from pytorch_optimizer.optimizer.avagrad import AvaGrad
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptSGD
@@ -61,7 +63,7 @@ from pytorch_optimizer.optimizer.ranger import Ranger
 from pytorch_optimizer.optimizer.ranger21 import Ranger21
 from pytorch_optimizer.optimizer.rotograd import RotoGrad
 from pytorch_optimizer.optimizer.sam import SAM
-from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD
+from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
 from pytorch_optimizer.optimizer.shampoo_utils import (
@@ -80,6 +82,7 @@ from pytorch_optimizer.optimizer.shampoo_utils import (
     power_iteration,
 )
 from pytorch_optimizer.optimizer.sm3 import SM3
+from pytorch_optimizer.optimizer.sophia import SophiaH
 from pytorch_optimizer.optimizer.srmm import SRMM
 from pytorch_optimizer.optimizer.swats import SWATS
 from pytorch_optimizer.optimizer.utils import (
@@ -145,6 +148,10 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     AvaGrad,
     AdaShift,
     AdaDelta,
+    Amos,
+    AdaHessian,
+    SophiaH,
+    SignSGD,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/base/optimizer.py

@@ -1,16 +1,84 @@
 import math
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
-from pytorch_optimizer.base.types import BETAS
+from pytorch_optimizer.base.types import BETAS, HUTCHINSON_G, PARAMETERS, STATE
 
 
 class BaseOptimizer(ABC):
     r"""Base optimizer class."""
 
+    @staticmethod
+    @torch.no_grad()
+    def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tensor]):
+        r"""Set hessian to state from external source. Generally useful when using functorch as a base.
+
+        Example usage:
+        ```
+        # Hutchinsons Estimator using HVP
+        noise = tree_map(lambda v: torch.randn_like(v), params)
+        loss_, hvp_est = jvp(grad(run_model_fn), (params,), (noise,))
+        hessian_diag_est  = tree_map(lambda a, b: a * b, hvp_est, noise)
+
+        optimizer.set_hessian(hessian_diag_est)
+        # OR
+        optimizer.step(hessian=hessian_diag_est)
+        ````
+        """
+        i: int = 0
+        for group in param_groups:
+            for p in group['params']:
+                if p.size() != hessian[i].size():
+                    raise ValueError(
+                        f'[-] the shape of parameter and hessian does not match. {p.size()} vs {hessian[i].size()}'
+                    )
+
+                state[p]['hessian'] = hessian[i]
+                i += 1
+
+    @staticmethod
+    @torch.no_grad()
+    def compute_hutchinson_hessian(
+        param_groups: PARAMETERS,
+        state: STATE,
+        num_samples: int = 1,
+        pre_zero: bool = True,
+        alpha: float = 1.0,
+        distribution: HUTCHINSON_G = 'gaussian',
+    ):
+        r"""Hutchinson's approximate hessian, added to the state under key `hessian`."""
+        if distribution not in ('gaussian', 'rademacher'):
+            raise NotImplementedError(f'[-] Hessian with distribution {distribution} is not implemented.')
+
+        params = []
+        for group in param_groups:
+            for p in group['params']:
+                if p.requires_grad and p.grad is not None and not p.grad.is_sparse:
+                    if 'hessian' not in state[p]:
+                        state[p]['hessian'] = torch.zeros_like(p)
+                    elif pre_zero:
+                        state[p]['hessian'].zero_()
+
+                    params.append(p)
+
+        if len(params) == 0:
+            return
+
+        grads = [p.grad for p in params]
+
+        for i in range(num_samples):
+            if distribution == 'rademacher':
+                zs = [torch.randint_like(p, 0, 1) * 2.0 - 1.0 for p in params]
+            else:
+                zs = [torch.randn_like(p) for p in params]
+
+            h_zs = torch.autograd.grad(grads, params, grad_outputs=zs, retain_graph=i < num_samples - 1)
+            for h_z, z, p in zip(h_zs, zs, params):
+                state[p]['hessian'].add_(h_z * z, alpha=(1 / num_samples) * alpha)
+
     @staticmethod
     def apply_weight_decay(
         p: torch.Tensor,
@@ -90,31 +158,6 @@ class BaseOptimizer(ABC):
 
         return grad * exp_grad_norm / grad_norm if exp_grad_norm > grad_norm else grad
 
-    @staticmethod
-    def validate_learning_rate(learning_rate: Optional[float]):
-        if learning_rate is not None and learning_rate < 0.0:
-            raise NegativeLRError(learning_rate)
-
-    @staticmethod
-    def validate_betas(betas: BETAS):
-        if not 0.0 <= betas[0] <= 1.0:
-            raise ValueError(f'[-] beta1 {betas[0]} must be in the range [0, 1]')
-        if not 0.0 <= betas[1] <= 1.0:
-            raise ValueError(f'[-] beta2 {betas[1]} must be in the range [0, 1]')
-
-        if len(betas) < 3:
-            return
-
-        if not 0.0 <= betas[2] <= 1.0:
-            raise ValueError(f'[-] beta3 {betas[2]} must be in the range [0, 1]')
-
-    @staticmethod
-    def validate_weight_decay_type(weight_decay_type: str):
-        if weight_decay_type not in ('l2', 'decoupled', 'stable'):
-            raise ValueError(
-                f'[-] weight_decay_type {weight_decay_type} must be one of (\'l2\', \'decoupled\', \'stable\')'
-            )
-
     @staticmethod
     def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)'):
         if range_type == '[)' and not low <= x < high:
@@ -149,37 +192,31 @@ class BaseOptimizer(ABC):
             raise NegativeStepError(step, step_type=step_type)
 
     @staticmethod
-    def validate_pullback_momentum(pullback_momentum: str):
-        if pullback_momentum not in ('none', 'reset', 'pullback'):
-            raise ValueError(
-                f'[-] pullback_momentum {pullback_momentum} must be one of (\'none\' or \'reset\' or \'pullback\')'
-            )
+    def validate_options(x: str, name: str, options: List[str]):
+        if x not in options:
+            opts: str = ' or '.join([f'\'{option}\'' for option in options]).strip()
+            raise ValueError(f'[-] {name} {x} must be one of ({opts})')
 
     @staticmethod
-    def validate_reduction(reduction: str):
-        if reduction not in ('mean', 'sum'):
-            raise ValueError(f'[-] reduction {reduction} must be one of (\'mean\' or \'sum\')')
+    def validate_learning_rate(learning_rate: Optional[float]):
+        if learning_rate is not None and learning_rate < 0.0:
+            raise NegativeLRError(learning_rate)
 
-    @staticmethod
-    def validate_rebound(rebound: str):
-        if rebound not in ('constant', 'belief'):
-            raise ValueError(f'[-] rebound {rebound} must be one of (\'constant\' or \'belief\')')
+    def validate_betas(self, betas: BETAS):
+        self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type='[]')
+        self.validate_range(betas[1], 'beta2', 0.0, 1.0, range_type='[]')
 
-    @staticmethod
-    def validate_a2grad_variant(variant: str):
-        if variant not in ('uni', 'inc', 'exp'):
-            raise ValueError(f'[-] A2Grad variant {variant} must be one of (\'uni\' or \'inc\' or \'exp\')')
+        if len(betas) < 3:
+            return
 
-    @staticmethod
-    def validate_nus(nus: Union[float, Tuple[float, float]]):
+        self.validate_range(betas[2], 'beta3', 0.0, 1.0, range_type='[]')
+
+    def validate_nus(self, nus: Union[float, Tuple[float, float]]):
         if isinstance(nus, float):
-            if not 0.0 <= nus <= 1.0:
-                raise ValueError(f'[-] nus {nus} must be in the range [0, 1]')
+            self.validate_range(nus, 'nu', 0.0, 1.0, range_type='[]')
         else:
-            if not 0.0 <= nus[0] <= 1.0:
-                raise ValueError(f'[-] nus1 {nus[0]} must be in the range [0, 1]')
-            if not 0.0 <= nus[1] <= 1.0:
-                raise ValueError(f'[-] nus2 {nus[1]} must be in the range [0, 1]')
+            self.validate_range(nus[0], 'nu1', 0.0, 1.0, range_type='[]')
+            self.validate_range(nus[1], 'nu2', 0.0, 1.0, range_type='[]')
 
     @abstractmethod
     def reset(self):

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/base/types.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple, Type, Union
 
 import torch
 from torch.optim import Optimizer
@@ -12,3 +12,5 @@ PARAMETERS = Optional[Union[Iterable[Dict[str, Any]], Iterable[torch.Tensor]]]
 STATE = Dict[str, Any]
 OPTIMIZER = Type[Optimizer]
 SCHEDULER = Type[_LRScheduler]
+
+HUTCHINSON_G = Literal['gaussian', 'rademacher']

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/a2grad.py

@@ -32,7 +32,7 @@ class A2Grad(Optimizer, BaseOptimizer):
         self.validate_learning_rate(lr)
         self.validate_non_negative(lips, 'lips')
         self.validate_non_negative(rho, 'rho')
-        self.validate_a2grad_variant(variant)
+        self.validate_options(variant, 'variant', ['uni', 'inc', 'exp'])
 
         self.variant = variant
 

pytorch_optimizer-2.10.0/pytorch_optimizer/optimizer/adahessian.py

@@ -0,0 +1,139 @@
+from typing import List, Optional
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, HUTCHINSON_G, LOSS, PARAMETERS
+
+
+class AdaHessian(Optimizer, BaseOptimizer):
+    r"""An Adaptive Second Order Optimizer for Machine Learning.
+
+        Requires `loss.backward(create_graph=True)` in order to calculate hessians.
+
+    :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 hessian_power: float. exponent of the hessian trace.
+    :param update_period: int. number of steps after which to apply hessian approximation.
+    :param num_samples: int. times to sample `z` for the approximation of the hessian trace.
+    :param hessian_distribution: HUTCHINSON_G. type of distribution to initialize hessian.
+    :param adam_debias: bool. Only correct the denominator to avoid inflating step sizes early in training.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-1,
+        betas: BETAS = (0.9, 0.999),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        hessian_power: float = 1.0,
+        update_period: int = 1,
+        num_samples: int = 1,
+        hessian_distribution: HUTCHINSON_G = 'rademacher',
+        adam_debias: bool = False,
+        eps: float = 1e-16,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+        self.validate_range(hessian_power, 'Hessian Power', 0, 1, range_type='(]')
+
+        self.update_period = update_period
+        self.num_samples = num_samples
+        self.distribution = hessian_distribution
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'hessian_power': hessian_power,
+            'adam_debias': adam_debias,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'AdaHessian'
+
+    @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_hessian_diag_sq'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None, hessian: Optional[List[torch.Tensor]] = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        step: int = self.param_groups[0].get('step', 1)
+
+        if hessian is not None:
+            self.set_hessian(self.param_groups, self.state, hessian)
+        elif step % self.update_period == 0:
+            self.compute_hutchinson_hessian(
+                self.param_groups, self.state, self.num_samples, distribution=self.distribution
+            )
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+
+            bias_correction1: float = 1.0 - beta1 ** group['step']
+            bias_correction2: float = 1.0 - beta2 ** group['step']
+
+            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 'exp_avg' not in state:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_hessian_diag_sq'] = 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_hessian_diag_sq = state['exp_avg'], state['exp_hessian_diag_sq']
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+
+                if 'hessian' in state and (group['step'] % self.update_period == 0 or hessian is not None):
+                    exp_hessian_diag_sq.mul_(beta2).addcmul_(state['hessian'], state['hessian'], value=1.0 - beta2)
+
+                de_nom = (exp_hessian_diag_sq / bias_correction2).pow_(group['hessian_power'] / 2).add_(group['eps'])
+
+                step_size: float = self.apply_adam_debias(group['adam_debias'], group['lr'], bias_correction1)
+                p.addcdiv_(exp_avg, de_nom, value=-step_size)
+
+        return loss

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/adai.py

@@ -103,7 +103,7 @@ class Adai(Optimizer, BaseOptimizer):
                 state['step'] += 1
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 bias_correction2: float = 1.0 - beta2 ** state['step']
 

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/adamp.py

@@ -122,7 +122,7 @@ class AdamP(Optimizer, BaseOptimizer):
                         state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 s_grad = self.get_adanorm_gradient(
                     grad=grad,

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/adan.py

@@ -130,7 +130,7 @@ class Adan(Optimizer, BaseOptimizer):
                 grad.mul_(clip_global_grad_norm)
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 grad_diff = state['previous_grad']
                 grad_diff.add_(grad)

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.10.0}/pytorch_optimizer/optimizer/agc.py

@@ -6,7 +6,7 @@ from pytorch_optimizer.optimizer.utils import unit_norm
 def agc(p: torch.Tensor, grad: torch.Tensor, agc_eps: float, agc_clip_val: float, eps: float = 1e-6) -> torch.Tensor:
     r"""Clip gradient values in excess of the unit wise norm.
 
-    :param p: Parameter. parameter.
+    :param p: torch.Tensor. parameter.
     :param grad: torch.Tensor, gradient.
     :param agc_eps: float. agc epsilon to clip the norm of parameter.
     :param agc_clip_val: float. norm clip.