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

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

@@ -1,15 +1,15 @@
 Metadata-Version: 2.1
 Name: pytorch_optimizer
-Version: 3.1.1
+Version: 3.2.0
 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,SOAP,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
 Maintainer-email: kozistr@gmail.com
-Requires-Python: >=3.8,<4.0.0
+Requires-Python: >=3.8
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Education
@@ -22,6 +22,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.6
 Classifier: Programming Language :: Python :: 3.7
 Classifier: Topic :: Scientific/Engineering
@@ -30,8 +31,9 @@ Classifier: Topic :: Software Development
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Provides-Extra: bitsandbytes
-Requires-Dist: bitsandbytes (>=0.43,<0.44) ; extra == "bitsandbytes"
-Requires-Dist: numpy ; python_version >= "3.8"
+Requires-Dist: bitsandbytes (>=0.44,<0.45) ; extra == "bitsandbytes"
+Requires-Dist: numpy (<=1.24.4) ; python_version < "3.9"
+Requires-Dist: numpy (>1.24.4) ; python_version >= "3.9"
 Requires-Dist: torch (>=1.10) ; python_version >= "3.8"
 Project-URL: Documentation, https://pytorch-optimizers.readthedocs.io/en/latest
 Project-URL: Repository, https://github.com/kozistr/pytorch_optimizer
@@ -49,7 +51,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, **77 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -66,16 +68,12 @@ So, please double-check the license before using it at your work.
 $ pip3 install pytorch-optimizer
 ```
 
-From `v2.12.0`, `v3.1.0`, you can use `bitsandbytes`, `q-galore-torch` optimizers respectively!
-please check [the bnb requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr), [q-galore-torch installation](https://github.com/VITA-Group/Q-GaLore?tab=readme-ov-file#install-q-galore-optimizer)
+From `v2.12.0`, `v3.1.0`, you can use `bitsandbytes`, `q-galore-torch`, `torchao` optimizers respectively!
+please check [the bnb requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr), [q-galore-torch installation](https://github.com/VITA-Group/Q-GaLore?tab=readme-ov-file#install-q-galore-optimizer), [torchao installation](https://github.com/pytorch/ao?tab=readme-ov-file#installation)
  before installing it.
 
 From `v3.0.0`, drop `Python 3.7` support. However, you can still use this package with `Python 3.7` by installing with `--ignore-requires-python` option.
 
-```bash
-$ pip install "pytorch-optimizer[bitsandbytes]"
-```
-
 ### Simple Usage
 
 ```python
@@ -133,6 +131,16 @@ from pytorch_optimizer import get_supported_optimizers
 supported_optimizers = get_supported_optimizers()
 ```
 
+or you can also search them with the filter(s).
+
+```python
+>>> get_supported_optimizers('adam*')
+['adamax', 'adamg', 'adammini', 'adamod', 'adamp', 'adams', 'adamw']
+
+>>> get_supported_optimizers(['adam*', 'ranger*'])
+['adamax', 'adamg', 'adammini', 'adamod', 'adamp', 'adams', 'adamw', 'ranger', 'ranger21']
+```
+
 | Optimizer     | Description                                                                                       | Official Code                                                                                                  | Paper                                                                                      | Citation                                                                                                                            |
 |---------------|---------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------|
 | AdaBelief     | *Adapting Step-sizes by the Belief in Observed Gradients*                                         | [github](https://github.com/juntang-zhuang/Adabelief-Optimizer)                                                | <https://arxiv.org/abs/2010.07468>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2020arXiv201007468Z/exportcitation)                                                        |
@@ -212,6 +220,8 @@ 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)                                         |
+| SOAP          | *Improving and Stabilizing Shampoo using Adam*                                                    | [github](https://github.com/nikhilvyas/SOAP)                                                                   | <https://arxiv.org/abs/2409.11321>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240911321V/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
@@ -223,6 +233,16 @@ from pytorch_optimizer import get_supported_lr_schedulers
 supported_lr_schedulers = get_supported_lr_schedulers()
 ```
 
+or you can also search them with the filter(s).
+
+```python
+>>> get_supported_lr_schedulers('cosine*')
+['cosine', 'cosine_annealing', 'cosine_annealing_with_warm_restart', 'cosine_annealing_with_warmup']
+
+>>> get_supported_lr_schedulers(['cosine*', '*warm*'])
+['cosine', 'cosine_annealing', 'cosine_annealing_with_warm_restart', 'cosine_annealing_with_warmup', 'warmup_stable_decay']
+```
+
 | LR Scheduler    | Description                                                                     | Official Code                                                                                                                       | Paper                              | Citation                                                                                           |
 |-----------------|---------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------|------------------------------------|----------------------------------------------------------------------------------------------------|
 | Explore-Exploit | *Wide-minima Density Hypothesis and the Explore-Exploit Learning Rate Schedule* |                                                                                                                                     | <https://arxiv.org/abs/2003.03977> | [cite](https://ui.adsabs.harvard.edu/abs/2020arXiv200303977I/exportcitation)                       |
@@ -240,6 +260,16 @@ from pytorch_optimizer import get_supported_loss_functions
 supported_loss_functions = get_supported_loss_functions()
 ```
 
+or you can also search them with the filter(s).
+
+```python
+>>> get_supported_loss_functions('*focal*')
+['bcefocalloss', 'focalcosineloss', 'focalloss', 'focaltverskyloss']
+
+>>> get_supported_loss_functions(['*focal*', 'bce*'])
+['bcefocalloss', 'bceloss', 'focalcosineloss', 'focalloss', 'focaltverskyloss']
+```
+
 | Loss Functions  | Description                                                                                                             | Official Code                                          | Paper                              | Citation                                                                     |
 |-----------------|-------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------|------------------------------------|------------------------------------------------------------------------------|
 | Label Smoothing | *Rethinking the Inception Architecture for Computer Vision*                                                             |                                                        | <https://arxiv.org/abs/1512.00567> | [cite](https://ui.adsabs.harvard.edu/abs/2015arXiv151200567S/exportcitation) |

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.2.0}/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, **77 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -27,16 +27,12 @@ So, please double-check the license before using it at your work.
 $ pip3 install pytorch-optimizer
 ```
 
-From `v2.12.0`, `v3.1.0`, you can use `bitsandbytes`, `q-galore-torch` optimizers respectively!
-please check [the bnb requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr), [q-galore-torch installation](https://github.com/VITA-Group/Q-GaLore?tab=readme-ov-file#install-q-galore-optimizer)
+From `v2.12.0`, `v3.1.0`, you can use `bitsandbytes`, `q-galore-torch`, `torchao` optimizers respectively!
+please check [the bnb requirements](https://github.com/TimDettmers/bitsandbytes?tab=readme-ov-file#tldr), [q-galore-torch installation](https://github.com/VITA-Group/Q-GaLore?tab=readme-ov-file#install-q-galore-optimizer), [torchao installation](https://github.com/pytorch/ao?tab=readme-ov-file#installation)
  before installing it.
 
 From `v3.0.0`, drop `Python 3.7` support. However, you can still use this package with `Python 3.7` by installing with `--ignore-requires-python` option.
 
-```bash
-$ pip install "pytorch-optimizer[bitsandbytes]"
-```
-
 ### Simple Usage
 
 ```python
@@ -94,6 +90,16 @@ from pytorch_optimizer import get_supported_optimizers
 supported_optimizers = get_supported_optimizers()
 ```
 
+or you can also search them with the filter(s).
+
+```python
+>>> get_supported_optimizers('adam*')
+['adamax', 'adamg', 'adammini', 'adamod', 'adamp', 'adams', 'adamw']
+
+>>> get_supported_optimizers(['adam*', 'ranger*'])
+['adamax', 'adamg', 'adammini', 'adamod', 'adamp', 'adams', 'adamw', 'ranger', 'ranger21']
+```
+
 | Optimizer     | Description                                                                                       | Official Code                                                                                                  | Paper                                                                                      | Citation                                                                                                                            |
 |---------------|---------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------|
 | AdaBelief     | *Adapting Step-sizes by the Belief in Observed Gradients*                                         | [github](https://github.com/juntang-zhuang/Adabelief-Optimizer)                                                | <https://arxiv.org/abs/2010.07468>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2020arXiv201007468Z/exportcitation)                                                        |
@@ -173,6 +179,8 @@ 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)                                         |
+| SOAP          | *Improving and Stabilizing Shampoo using Adam*                                                    | [github](https://github.com/nikhilvyas/SOAP)                                                                   | <https://arxiv.org/abs/2409.11321>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240911321V/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 
@@ -184,6 +192,16 @@ from pytorch_optimizer import get_supported_lr_schedulers
 supported_lr_schedulers = get_supported_lr_schedulers()
 ```
 
+or you can also search them with the filter(s).
+
+```python
+>>> get_supported_lr_schedulers('cosine*')
+['cosine', 'cosine_annealing', 'cosine_annealing_with_warm_restart', 'cosine_annealing_with_warmup']
+
+>>> get_supported_lr_schedulers(['cosine*', '*warm*'])
+['cosine', 'cosine_annealing', 'cosine_annealing_with_warm_restart', 'cosine_annealing_with_warmup', 'warmup_stable_decay']
+```
+
 | LR Scheduler    | Description                                                                     | Official Code                                                                                                                       | Paper                              | Citation                                                                                           |
 |-----------------|---------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------|------------------------------------|----------------------------------------------------------------------------------------------------|
 | Explore-Exploit | *Wide-minima Density Hypothesis and the Explore-Exploit Learning Rate Schedule* |                                                                                                                                     | <https://arxiv.org/abs/2003.03977> | [cite](https://ui.adsabs.harvard.edu/abs/2020arXiv200303977I/exportcitation)                       |
@@ -201,6 +219,16 @@ from pytorch_optimizer import get_supported_loss_functions
 supported_loss_functions = get_supported_loss_functions()
 ```
 
+or you can also search them with the filter(s).
+
+```python
+>>> get_supported_loss_functions('*focal*')
+['bcefocalloss', 'focalcosineloss', 'focalloss', 'focaltverskyloss']
+
+>>> get_supported_loss_functions(['*focal*', 'bce*'])
+['bcefocalloss', 'bceloss', 'focalcosineloss', 'focalloss', 'focaltverskyloss']
+```
+
 | Loss Functions  | Description                                                                                                             | Official Code                                          | Paper                              | Citation                                                                     |
 |-----------------|-------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------|------------------------------------|------------------------------------------------------------------------------|
 | Label Smoothing | *Rethinking the Inception Architecture for Computer Vision*                                                             |                                                        | <https://arxiv.org/abs/1512.00567> | [cite](https://ui.adsabs.harvard.edu/abs/2015arXiv151200567S/exportcitation) |

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.1.1"
+version = "3.2.0"
 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", "SOAP", "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",
@@ -35,6 +35,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Operating System :: OS Independent",
     "Topic :: Scientific/Engineering",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
@@ -44,10 +45,13 @@ classifiers = [
 ]
 
 [tool.poetry.dependencies]
-python = ">=3.8,<4.0.0"
-numpy = { version = "*", python = ">=3.8" }
+python = ">=3.8"
+numpy = [
+    { version = ">1.24.4", python = ">=3.9" },
+    { version = "<=1.24.4", python = "<3.9" },
+]
 torch = { version = ">=1.10", python = ">=3.8", source = "torch" }
-bitsandbytes = { version = "^0.43", optional = true }
+bitsandbytes = { version = "^0.44", optional = true }
 
 [tool.poetry.dev-dependencies]
 isort = { version = "^5", python = ">=3.8" }
@@ -65,20 +69,13 @@ url = "https://download.pytorch.org/whl/cpu"
 priority = "explicit"
 
 [tool.ruff]
-lint.select = [
-    "A", "B", "C4", "D", "E", "F", "G", "I", "N", "S", "T", "ISC", "ICN", "W", "INP", "PIE", "T20", "RET", "SIM",
-    "TID", "ARG", "ERA", "RUF", "YTT", "PL", "Q"
-]
-lint.ignore = [
-    "B905", "D100", "D102", "D104", "D105", "D107", "D203", "D213", "D413", "PIE790", "PLR0912", "PLR0913", "PLR0915",
-    "PLR2004", "RUF013", "Q003",
+src = [
+    "pytorch_optimizer",
+    "tests",
+    "examples",
 ]
-lint.fixable = ["ALL"]
-lint.unfixable = ["F401"]
-lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-lint.flake8-quotes.docstring-quotes = "double"
-lint.flake8-quotes.inline-quotes = "single"
-lint.pylint.max-args = 7
+target-version = "py312"
+line-length = 119
 exclude = [
     ".git",
     ".github",
@@ -92,32 +89,40 @@ exclude = [
     ".venv",
     "__pypackages__",
 ]
-line-length = 119
-target-version = "py311"
+
+[tool.ruff.lint]
+select = [
+    "A", "B", "C4", "D", "E", "F", "G", "I", "N", "S", "T", "ISC", "ICN", "W", "INP", "PIE", "T20", "RET", "SIM",
+    "TID", "ARG", "ERA", "RUF", "YTT", "PL", "Q"
+]
+ignore = [
+    "B905", "D100", "D102", "D104", "D105", "D107", "D203", "D213", "D413", "PIE790", "PLR0912", "PLR0913", "PLR0915",
+    "PLR2004", "RUF013", "Q003", "ARG002",
+]
+fixable = ["ALL"]
+unfixable = ["F401"]
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+flake8-quotes.docstring-quotes = "double"
+flake8-quotes.inline-quotes = "single"
+
+[tool.ruff.lint.extend-per-file-ignores]
+"hubconf.py" = ["D", "INP001"]
+"examples/visualize_optimizers.py" = ["D103", "D400", "D415"]
+"**/__init__.py" = ["F401"]
+"{tests,examples}/*.py" = ["D", "S101"]
+
+[tool.ruff.lint.isort]
+combine-as-imports = false
+detect-same-package = true
+force-sort-within-sections = false
+known-first-party = ["pytorch_optimizer"]
+
+[tool.ruff.lint.pylint]
+max-args = 7
 
 [tool.ruff.format]
 quote-style = "single"
 
-[tool.ruff.lint.per-file-ignores]
-"./pytorch_optimizer/__init__.py" = ["F401"]
-"./pytorch_optimizer/lr_scheduler/__init__.py" = ["F401"]
-"./hubconf.py" = ["D", "INP001"]
-"./examples/visualize_optimizers.py" = ["D103", "D400", "D415"]
-"./tests/__init__.py" = ["D"]
-"./tests/constants.py" = ["D"]
-"./tests/utils.py" = ["D"]
-"./tests/test_base.py" = ["D", "S101"]
-"./tests/test_utils.py" = ["D", "S101", "ERA001"]
-"./tests/test_gradients.py" = ["D", "S101"]
-"./tests/test_optimizers.py" = ["D", "S101"]
-"./tests/test_optimizer_parameters.py" = ["D", "S101"]
-"./tests/test_general_optimizer_parameters.py" = ["D", "S101"]
-"./tests/test_lr_schedulers.py" = ["D", "S101"]
-"./tests/test_lr_scheduler_parameters.py" = ["D", "S101"]
-"./tests/test_create_optimizer.py" = ["D"]
-"./tests/test_loss_functions.py" = ["D", "S101"]
-"./tests/test_load_modules.py" = ["D", "S101"]
-
 [tool.pytest.ini_options]
 testpaths = "tests"
 
@@ -128,5 +133,5 @@ omit = [
 ]
 
 [build-system]
-requires = ["poetry-core>=1.0.0"]
+requires = ["poetry-core>=1.4.0"]
 build-backend = "poetry.core.masonry.api"

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

@@ -1,6 +1,7 @@
 # ruff: noqa
+import fnmatch
 from importlib.util import find_spec
-from typing import Dict, List
+from typing import Dict, List, Optional, Sequence, Set, Union
 
 import torch.cuda
 from torch import nn
@@ -54,6 +55,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
@@ -113,12 +115,14 @@ from pytorch_optimizer.optimizer.shampoo_utils import (
     power_iteration,
 )
 from pytorch_optimizer.optimizer.sm3 import SM3
+from pytorch_optimizer.optimizer.soap import SOAP
 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.tiger import Tiger
 from pytorch_optimizer.optimizer.trac import TRAC
 from pytorch_optimizer.optimizer.utils import (
+    CPUOffloadOptimizer,
     clip_grad_norm,
     disable_running_stats,
     enable_running_stats,
@@ -132,6 +136,7 @@ from pytorch_optimizer.optimizer.yogi import Yogi
 
 HAS_BNB: bool = find_spec('bitsandbytes') is not None
 HAS_Q_GALORE: bool = find_spec('q-galore-torch') is not None
+HAS_TORCHAO: bool = find_spec('torchao') is not None
 
 OPTIMIZER_LIST: List[OPTIMIZER] = [
     AdamW,
@@ -208,6 +213,8 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     AdamMini,
     AdaLOMO,
     AdamG,
+    AdEMAMix,
+    SOAP,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
@@ -297,6 +304,15 @@ def load_bnb_optimizer(optimizer: str) -> OPTIMIZER:  # pragma: no cover
         return optim.RMSprop32bit
     if 'sgd32bit' in optimizer:
         return optim.SGD32bit
+    if 'ademamix8bit' in optimizer:
+        return optim.AdEMAMix8bit
+    if 'ademamix32bit' in optimizer:
+        return optim.AdEMAMix32bit
+    if 'paged_ademamix8bit' in optimizer:
+        return optim.PagedAdEMAMix8bit
+    if 'paged_ademamix32bit' in optimizer:
+        return optim.PagedAdEMAMix32bit
+
     raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
 
 
@@ -310,19 +326,40 @@ def load_q_galore_optimizer(optimizer: str) -> OPTIMIZER:  # pragma: no cover
     raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
 
 
+def load_ao_optimizer(optimizer: str) -> OPTIMIZER:  # pragma: no cover
+    r"""load TorchAO optimizer instance."""
+    from torchao.prototype import low_bit_optim
+
+    if 'adamw8bit' in optimizer:
+        return low_bit_optim.AdamW8bit
+    if 'adamw4bit' in optimizer:
+        return low_bit_optim.AdamW4bit
+    if 'adamwfp8' in optimizer:
+        return low_bit_optim.AdamWFp8
+
+    raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
+
+
 def load_optimizer(optimizer: str) -> OPTIMIZER:
     optimizer: str = optimizer.lower()
 
     if optimizer.startswith('bnb'):
         if HAS_BNB and torch.cuda.is_available():
             return load_bnb_optimizer(optimizer)  # pragma: no cover
-        raise ImportError(f'[-] bitsandbytes and CUDA required for the optimizer {optimizer}')
+        raise ImportError(f'bitsandbytes and CUDA required for the optimizer {optimizer}')
     if optimizer.startswith('q_galore'):
         if HAS_Q_GALORE and torch.cuda.is_available():
             return load_q_galore_optimizer(optimizer)  # pragma: no cover
-        raise ImportError(f'[-] bitsandbytes, q-galore-torch, and CUDA required for the optimizer {optimizer}')
+        raise ImportError(f'bitsandbytes, q-galore-torch, and CUDA required for the optimizer {optimizer}')
+    if optimizer.startswith('torchao'):
+        if HAS_TORCHAO and torch.cuda.is_available():
+            return load_ao_optimizer(optimizer)  # pragma: no cover
+        raise ImportError(
+            f'torchao required for the optimizer {optimizer}. '
+            'usage: https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#usage'
+        )
     if optimizer not in OPTIMIZERS:
-        raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
+        raise NotImplementedError(f'not implemented optimizer : {optimizer}')
 
     return OPTIMIZERS[optimizer]
 
@@ -380,13 +417,55 @@ def load_lr_scheduler(lr_scheduler: str) -> SCHEDULER:
     return LR_SCHEDULERS[lr_scheduler]
 
 
-def get_supported_optimizers() -> List[OPTIMIZER]:
-    return OPTIMIZER_LIST
+def get_supported_optimizers(filters: Optional[Union[str, List[str]]] = None) -> List[str]:
+    r"""Return list of available optimizer names, sorted alphabetically.
+
+    :param filters: Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will
+        return the whole list.
+    """
+    if filters is None:
+        return sorted(OPTIMIZERS.keys())
+
+    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]
+
+    filtered_list: Set[str] = set()
+    for include_filter in include_filters:
+        filtered_list.update(fnmatch.filter(OPTIMIZERS.keys(), include_filter))
 
+    return sorted(filtered_list)
+
+
+def get_supported_lr_schedulers(filters: Optional[Union[str, List[str]]] = None) -> List[str]:
+    r"""Return list of available lr scheduler names, sorted alphabetically.
+
+    :param filters: Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will
+        return the whole list.
+    """
+    if filters is None:
+        return sorted(LR_SCHEDULERS.keys())
+
+    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]
+
+    filtered_list: Set[str] = set()
+    for include_filter in include_filters:
+        filtered_list.update(fnmatch.filter(LR_SCHEDULERS.keys(), include_filter))
+
+    return sorted(filtered_list)
+
+
+def get_supported_loss_functions(filters: Optional[Union[str, List[str]]] = None) -> List[str]:
+    r"""Return list of available loss function names, sorted alphabetically.
+
+    :param filters: Optional[Union[str, List[str]]]. wildcard filter string that works with fmatch. if None, it will
+        return the whole list.
+    """
+    if filters is None:
+        return sorted(LOSS_FUNCTIONS.keys())
 
-def get_supported_lr_schedulers() -> List[SCHEDULER]:
-    return list(LR_SCHEDULER_LIST.values())
+    include_filters: Sequence[str] = filters if isinstance(filters, (tuple, list)) else [filters]
 
+    filtered_list: Set[str] = set()
+    for include_filter in include_filters:
+        filtered_list.update(fnmatch.filter(LOSS_FUNCTIONS.keys(), include_filter))
 
-def get_supported_loss_functions() -> List[nn.Module]:
-    return LOSS_FUNCTION_LIST
+    return sorted(filtered_list)

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.2.0}/pytorch_optimizer/base/types.py

@@ -15,3 +15,5 @@ SCHEDULER = Type[LRScheduler]
 
 HUTCHINSON_G = Literal['gaussian', 'rademacher']
 CLASS_MODE = Literal['binary', 'multiclass', 'multilabel']
+
+DATA_FORMAT = Literal['channels_first', 'channels_last']

{pytorch_optimizer-3.1.1 → pytorch_optimizer-3.2.0}/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.2.0}/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.2.0}/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.2.0}/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.2.0}/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.2.0}/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.2.0}/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.2.0}/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.2.0}/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)