Repository: kozistr/pytorch_optimizer

Issue:
Support SOAP
https://github.com/nikhilvyas/SOAP

Hints from issue comments:

Solution patch:
diff --git a/README.md b/README.md
index acd0f1bde..2e1b77e87 100644
--- a/README.md
+++ b/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, **76 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **77 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -174,6 +174,7 @@ supported_optimizers = get_supported_optimizers()
 | 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
 
diff --git a/docs/changelogs/v3.2.0.md b/docs/changelogs/v3.2.0.md
new file mode 100644
index 000000000..71aa85cd0
--- /dev/null
+++ b/docs/changelogs/v3.2.0.md
@@ -0,0 +1,6 @@
+## Change Log
+
+### Feature
+
+* Implement `SOAP` optimizer. (#275)
+    * [SOAP: Improving and Stabilizing Shampoo using Adam](https://arxiv.org/abs/2409.11321) 
diff --git a/docs/index.md b/docs/index.md
index acd0f1bde..2e1b77e87 100644
--- a/docs/index.md
+++ b/docs/index.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, **76 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **77 optimizers (+ `bitsandbytes`, `qgalore`)**, **16 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -174,6 +174,7 @@ supported_optimizers = get_supported_optimizers()
 | 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
 
diff --git a/docs/optimizer.md b/docs/optimizer.md
index 35f1d2939..3fbf42862 100644
--- a/docs/optimizer.md
+++ b/docs/optimizer.md
@@ -320,6 +320,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.SOAP
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.SophiaH
     :docstring:
     :members:
diff --git a/poetry.lock b/poetry.lock
index 9d9f4079e..79d74596f 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -193,18 +193,18 @@ test = ["pytest (>=6)"]
 
 [[package]]
 name = "filelock"
-version = "3.16.0"
+version = "3.16.1"
 description = "A platform independent file lock."
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "filelock-3.16.0-py3-none-any.whl", hash = "sha256:f6ed4c963184f4c84dd5557ce8fece759a3724b37b80c6c4f20a2f63a4dc6609"},
-    {file = "filelock-3.16.0.tar.gz", hash = "sha256:81de9eb8453c769b63369f87f11131a7ab04e367f8d97ad39dc230daa07e3bec"},
+    {file = "filelock-3.16.1-py3-none-any.whl", hash = "sha256:2082e5703d51fbf98ea75855d9d5527e33d8ff23099bec374a134febee6946b0"},
+    {file = "filelock-3.16.1.tar.gz", hash = "sha256:c249fbfcd5db47e5e2d6d62198e565475ee65e4831e2561c8e313fa7eb961435"},
 ]
 
 [package.extras]
-docs = ["furo (>=2024.8.6)", "sphinx (>=8.0.2)", "sphinx-autodoc-typehints (>=2.4)"]
-testing = ["covdefaults (>=2.3)", "coverage (>=7.6.1)", "diff-cover (>=9.1.1)", "pytest (>=8.3.2)", "pytest-asyncio (>=0.24)", "pytest-cov (>=5)", "pytest-mock (>=3.14)", "pytest-timeout (>=2.3.1)", "virtualenv (>=20.26.3)"]
+docs = ["furo (>=2024.8.6)", "sphinx (>=8.0.2)", "sphinx-autodoc-typehints (>=2.4.1)"]
+testing = ["covdefaults (>=2.3)", "coverage (>=7.6.1)", "diff-cover (>=9.2)", "pytest (>=8.3.3)", "pytest-asyncio (>=0.24)", "pytest-cov (>=5)", "pytest-mock (>=3.14)", "pytest-timeout (>=2.3.1)", "virtualenv (>=20.26.4)"]
 typing = ["typing-extensions (>=4.12.2)"]
 
 [[package]]
@@ -464,13 +464,13 @@ files = [
 
 [[package]]
 name = "platformdirs"
-version = "4.3.2"
+version = "4.3.6"
 description = "A small Python package for determining appropriate platform-specific dirs, e.g. a `user data dir`."
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "platformdirs-4.3.2-py3-none-any.whl", hash = "sha256:eb1c8582560b34ed4ba105009a4badf7f6f85768b30126f351328507b2beb617"},
-    {file = "platformdirs-4.3.2.tar.gz", hash = "sha256:9e5e27a08aa095dd127b9f2e764d74254f482fef22b0970773bfba79d091ab8c"},
+    {file = "platformdirs-4.3.6-py3-none-any.whl", hash = "sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb"},
+    {file = "platformdirs-4.3.6.tar.gz", hash = "sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907"},
 ]
 
 [package.extras]
@@ -495,13 +495,13 @@ testing = ["pytest", "pytest-benchmark"]
 
 [[package]]
 name = "pytest"
-version = "8.3.2"
+version = "8.3.3"
 description = "pytest: simple powerful testing with Python"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "pytest-8.3.2-py3-none-any.whl", hash = "sha256:4ba08f9ae7dcf84ded419494d229b48d0903ea6407b030eaec46df5e6a73bba5"},
-    {file = "pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce"},
+    {file = "pytest-8.3.3-py3-none-any.whl", hash = "sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2"},
+    {file = "pytest-8.3.3.tar.gz", hash = "sha256:70b98107bd648308a7952b06e6ca9a50bc660be218d53c257cc1fc94fda10181"},
 ]
 
 [package.dependencies]
@@ -535,40 +535,40 @@ testing = ["fields", "hunter", "process-tests", "pytest-xdist", "virtualenv"]
 
 [[package]]
 name = "ruff"
-version = "0.6.4"
+version = "0.6.6"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "ruff-0.6.4-py3-none-linux_armv6l.whl", hash = "sha256:c4b153fc152af51855458e79e835fb6b933032921756cec9af7d0ba2aa01a258"},
-    {file = "ruff-0.6.4-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:bedff9e4f004dad5f7f76a9d39c4ca98af526c9b1695068198b3bda8c085ef60"},
-    {file = "ruff-0.6.4-py3-none-macosx_11_0_arm64.whl", hash = "sha256:d02a4127a86de23002e694d7ff19f905c51e338c72d8e09b56bfb60e1681724f"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7862f42fc1a4aca1ea3ffe8a11f67819d183a5693b228f0bb3a531f5e40336fc"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:eebe4ff1967c838a1a9618a5a59a3b0a00406f8d7eefee97c70411fefc353617"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:932063a03bac394866683e15710c25b8690ccdca1cf192b9a98260332ca93408"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:50e30b437cebef547bd5c3edf9ce81343e5dd7c737cb36ccb4fe83573f3d392e"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c44536df7b93a587de690e124b89bd47306fddd59398a0fb12afd6133c7b3818"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ea086601b22dc5e7693a78f3fcfc460cceabfdf3bdc36dc898792aba48fbad6"},
-    {file = "ruff-0.6.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b52387d3289ccd227b62102c24714ed75fbba0b16ecc69a923a37e3b5e0aaaa"},
-    {file = "ruff-0.6.4-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:0308610470fcc82969082fc83c76c0d362f562e2f0cdab0586516f03a4e06ec6"},
-    {file = "ruff-0.6.4-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:803b96dea21795a6c9d5bfa9e96127cc9c31a1987802ca68f35e5c95aed3fc0d"},
-    {file = "ruff-0.6.4-py3-none-musllinux_1_2_i686.whl", hash = "sha256:66dbfea86b663baab8fcae56c59f190caba9398df1488164e2df53e216248baa"},
-    {file = "ruff-0.6.4-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:34d5efad480193c046c86608dbba2bccdc1c5fd11950fb271f8086e0c763a5d1"},
-    {file = "ruff-0.6.4-py3-none-win32.whl", hash = "sha256:f0f8968feea5ce3777c0d8365653d5e91c40c31a81d95824ba61d871a11b8523"},
-    {file = "ruff-0.6.4-py3-none-win_amd64.whl", hash = "sha256:549daccee5227282289390b0222d0fbee0275d1db6d514550d65420053021a58"},
-    {file = "ruff-0.6.4-py3-none-win_arm64.whl", hash = "sha256:ac4b75e898ed189b3708c9ab3fc70b79a433219e1e87193b4f2b77251d058d14"},
-    {file = "ruff-0.6.4.tar.gz", hash = "sha256:ac3b5bfbee99973f80aa1b7cbd1c9cbce200883bdd067300c22a6cc1c7fba212"},
+    {file = "ruff-0.6.6-py3-none-linux_armv6l.whl", hash = "sha256:f5bc5398457484fc0374425b43b030e4668ed4d2da8ee7fdda0e926c9f11ccfb"},
+    {file = "ruff-0.6.6-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:515a698254c9c47bb84335281a170213b3ee5eb47feebe903e1be10087a167ce"},
+    {file = "ruff-0.6.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:6bb1b4995775f1837ab70f26698dd73852bbb82e8f70b175d2713c0354fe9182"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69c546f412dfae8bb9cc4f27f0e45cdd554e42fecbb34f03312b93368e1cd0a6"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:59627e97364329e4eae7d86fa7980c10e2b129e2293d25c478ebcb861b3e3fd6"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:94c3f78c3d32190aafbb6bc5410c96cfed0a88aadb49c3f852bbc2aa9783a7d8"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:704da526c1e137f38c8a067a4a975fe6834b9f8ba7dbc5fd7503d58148851b8f"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:efeede5815a24104579a0f6320660536c5ffc1c91ae94f8c65659af915fb9de9"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e368aef0cc02ca3593eae2fb8186b81c9c2b3f39acaaa1108eb6b4d04617e61f"},
+    {file = "ruff-0.6.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2653fc3b2a9315bd809725c88dd2446550099728d077a04191febb5ea79a4f79"},
+    {file = "ruff-0.6.6-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:bb858cd9ce2d062503337c5b9784d7b583bcf9d1a43c4df6ccb5eab774fbafcb"},
+    {file = "ruff-0.6.6-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:488f8e15c01ea9afb8c0ba35d55bd951f484d0c1b7c5fd746ce3c47ccdedce68"},
+    {file = "ruff-0.6.6-py3-none-musllinux_1_2_i686.whl", hash = "sha256:aefb0bd15f1cfa4c9c227b6120573bb3d6c4ee3b29fb54a5ad58f03859bc43c6"},
+    {file = "ruff-0.6.6-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a4c0698cc780bcb2c61496cbd56b6a3ac0ad858c966652f7dbf4ceb029252fbe"},
+    {file = "ruff-0.6.6-py3-none-win32.whl", hash = "sha256:aadf81ddc8ab5b62da7aae78a91ec933cbae9f8f1663ec0325dae2c364e4ad84"},
+    {file = "ruff-0.6.6-py3-none-win_amd64.whl", hash = "sha256:0adb801771bc1f1b8cf4e0a6fdc30776e7c1894810ff3b344e50da82ef50eeb1"},
+    {file = "ruff-0.6.6-py3-none-win_arm64.whl", hash = "sha256:4b4d32c137bc781c298964dd4e52f07d6f7d57c03eae97a72d97856844aa510a"},
+    {file = "ruff-0.6.6.tar.gz", hash = "sha256:0fc030b6fd14814d69ac0196396f6761921bd20831725c7361e1b8100b818034"},
 ]
 
 [[package]]
 name = "sympy"
-version = "1.13.2"
+version = "1.13.3"
 description = "Computer algebra system (CAS) in Python"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "sympy-1.13.2-py3-none-any.whl", hash = "sha256:c51d75517712f1aed280d4ce58506a4a88d635d6b5dd48b39102a7ae1f3fcfe9"},
-    {file = "sympy-1.13.2.tar.gz", hash = "sha256:401449d84d07be9d0c7a46a64bd54fe097667d5e7181bfe67ec777be9e01cb13"},
+    {file = "sympy-1.13.3-py3-none-any.whl", hash = "sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73"},
+    {file = "sympy-1.13.3.tar.gz", hash = "sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9"},
 ]
 
 [package.dependencies]
diff --git a/pyproject.toml b/pyproject.toml
index c214a2e9d..7d2074313 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,9 +17,9 @@ keywords = [
     "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",
+    "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",
diff --git a/pytorch_optimizer/__init__.py b/pytorch_optimizer/__init__.py
index c9279b17f..3ab9cc9f4 100644
--- a/pytorch_optimizer/__init__.py
+++ b/pytorch_optimizer/__init__.py
@@ -114,6 +114,7 @@
     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
@@ -210,6 +211,7 @@
     AdaLOMO,
     AdamG,
     AdEMAMix,
+    SOAP,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
diff --git a/pytorch_optimizer/base/types.py b/pytorch_optimizer/base/types.py
index 422ad1f73..ff6aaaf9c 100644
--- a/pytorch_optimizer/base/types.py
+++ b/pytorch_optimizer/base/types.py
@@ -15,3 +15,5 @@
 
 HUTCHINSON_G = Literal['gaussian', 'rademacher']
 CLASS_MODE = Literal['binary', 'multiclass', 'multilabel']
+
+DATA_FORMAT = Literal['channels_first', 'channels_last']
diff --git a/pytorch_optimizer/optimizer/soap.py b/pytorch_optimizer/optimizer/soap.py
new file mode 100644
index 000000000..792fd6fb6
--- /dev/null
+++ b/pytorch_optimizer/optimizer/soap.py
@@ -0,0 +1,335 @@
+import math
+from itertools import chain
+from typing import List, 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, DATA_FORMAT, DEFAULTS, LOSS, PARAMETERS
+from pytorch_optimizer.optimizer.shampoo_utils import merge_small_dims
+
+
+class SOAP(BaseOptimizer):
+    r"""Improving and Stabilizing Shampoo using Adam.
+
+    :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 shampoo_beta: Optional[float]. if not None, use this beta for the pre-conditioner (L and R in paper,
+        state['GG'] below) moving average instead of betas[1].
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param precondition_frequency: int. how often to update the pre-conditioner.
+    :param max_precondition_dim: int. maximum dimension of the pre-conditioner. Set to 10000, so that we exclude most
+        common vocab sizes while including layers.
+    :param merge_dims: bool. whether to merge dimensions of the pre-conditioner
+    :param precondition_1d: bool. whether to precondition 1D gradients.
+    :param correct_bias: bool. whether to correct bias in Adam.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 3e-3,
+        betas: BETAS = (0.95, 0.95),
+        shampoo_beta: Optional[float] = None,
+        weight_decay: float = 1e-2,
+        precondition_frequency: int = 10,
+        max_precondition_dim: int = 10000,
+        merge_dims: bool = False,
+        precondition_1d: bool = False,
+        correct_bias: bool = True,
+        data_format: DATA_FORMAT = 'channels_first',
+        eps: float = 1e-8,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(shampoo_beta, 'shampoo_beta')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_positive(precondition_frequency, 'precondition_frequency')
+        self.validate_positive(max_precondition_dim, 'max_precondition_dim')
+        self.validate_non_negative(eps, 'eps')
+
+        self.data_format = data_format
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'shampoo_beta': shampoo_beta,
+            'weight_decay': weight_decay,
+            'precondition_frequency': precondition_frequency,
+            'max_precondition_dim': max_precondition_dim,
+            'merge_dims': merge_dims,
+            'precondition_1d': precondition_1d,
+            'correct_bias': correct_bias,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'SOAP'
+
+    @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)
+
+    def project(
+        self,
+        grad: torch.Tensor,
+        state,
+        merge_dims: bool = False,
+        max_precondition_dim: int = 10000,
+        project_type: str = 'forward',
+    ) -> torch.Tensor:
+        original_shape = grad.shape
+
+        if merge_dims:
+            if self.data_format == 'channels_last' and grad.dim() == 4:
+                permuted_shape = grad.permute(0, 3, 1, 2).shape
+
+            grad = grad.reshape(merge_small_dims(grad.size(), max_precondition_dim))
+
+        for mat in state['Q']:
+            if len(mat) > 0:
+                grad = torch.tensordot(grad, mat, dims=[[0], [0 if project_type == 'forward' else 1]])
+            else:
+                grad = grad.permute([*list(range(1, len(grad.shape))), 0])
+
+        if merge_dims:
+            if self.data_format == 'channels_last' and len(original_shape) == 4:
+                grad = grad.reshape(permuted_shape).permute(0, 2, 3, 1)
+            else:
+                grad = grad.reshape(original_shape)
+
+        return grad
+
+    @staticmethod
+    def get_orthogonal_matrix(mat: torch.Tensor) -> List[torch.Tensor]:
+        matrices: List = []
+        for m in mat:
+            if len(m) == 0:
+                matrices.append([])
+                continue
+
+            try:
+                _, q = torch.linalg.eigh(m + 1e-30 * torch.eye(m.shape[0], device=m.device))
+            except Exception:  # pragma: no cover
+                _, q = torch.linalg.eigh(m.to(torch.float64) + 1e-30 * torch.eye(m.shape[0], device=m.device))
+                q = q.to(m.dtype)
+
+            q = torch.flip(q, dims=[1])
+
+            matrices.append(q)
+
+        return matrices
+
+    def get_orthogonal_matrix_qr(self, state, max_precondition_dim: int = 10000, merge_dims: bool = False):
+        r"""Compute the eigen-bases of the pre-conditioner using one round of power iteration."""
+        orig_shape = state['exp_avg_sq'].shape
+        if self.data_format == 'channels_last' and len(orig_shape) == 4:
+            permuted_shape = state['exp_avg_sq'].permute(0, 3, 1, 2).shape
+
+        exp_avg_sq = state['exp_avg_sq']
+        if merge_dims:
+            exp_avg_sq = exp_avg_sq.reshape(merge_small_dims(exp_avg_sq.size(), max_precondition_dim))
+
+        matrices = []
+        for ind, (m, o) in enumerate(zip(state['GG'], state['Q'])):
+            if len(m) == 0:
+                matrices.append([])
+                continue
+
+            est_eig = torch.diag(o.T @ m @ o)
+            sort_idx = torch.argsort(est_eig, descending=True)
+            exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
+
+            power_iter = m @ o[:, sort_idx]
+
+            q, _ = torch.linalg.qr(power_iter)
+
+            matrices.append(q)
+
+        if merge_dims:
+            if self.data_format == 'channels_last' and len(orig_shape) == 4:
+                exp_avg_sq = exp_avg_sq.reshape(permuted_shape).permute(0, 2, 3, 1)
+            else:
+                exp_avg_sq = exp_avg_sq.reshape(orig_shape)
+
+        state['exp_avg_sq'] = exp_avg_sq
+
+        return matrices
+
+    @staticmethod
+    def init_pre_conditioner(
+        grad,
+        state,
+        precondition_frequency: int = 10,
+        shampoo_beta: float = 0.95,
+        max_precondition_dim: int = 10000,
+        precondition_1d: bool = False,
+        merge_dims: bool = False,
+    ) -> None:
+        state['GG'] = []
+        if grad.dim() == 1:
+            if not precondition_1d or grad.shape[0] > max_precondition_dim:
+                state['GG'].append([])
+            else:
+                state['GG'].append(torch.zeros(grad.shape[0], grad.shape[0], device=grad.device))
+        else:
+            if merge_dims:
+                grad = grad.reshape(merge_small_dims(grad.size(), max_precondition_dim))
+
+            for sh in grad.shape:
+                if sh > max_precondition_dim:
+                    state['GG'].append([])
+                else:
+                    state['GG'].append(torch.zeros(sh, sh, device=grad.device))
+
+        state['Q'] = None
+        state['precondition_frequency'] = precondition_frequency
+        state['shampoo_beta'] = shampoo_beta
+
+    def update_pre_conditioner(
+        self,
+        grad,
+        state,
+        step: int,
+        max_precondition_dim: int = 10000,
+        precondition_1d: bool = False,
+        merge_dims: bool = False,
+    ) -> None:
+        if grad.dim() == 1:
+            if precondition_1d and grad.shape[0] <= max_precondition_dim:
+                state['GG'][0].lerp_(
+                    (grad.unsqueeze(1) @ grad.unsqueeze(0)).to(state['GG'][0].dtype),
+                    weight=1.0 - state['shampoo_beta'],
+                )
+        else:
+            if merge_dims:
+                grad = grad.reshape(merge_small_dims(grad.size(), max_precondition_dim))
+
+            for idx, dim in enumerate(grad.shape):
+                if dim <= max_precondition_dim:
+                    outer_product = torch.tensordot(
+                        grad,
+                        grad,
+                        dims=[[*chain(range(idx), range(idx + 1, len(grad.shape)))]] * 2,
+                    )
+
+                    state['GG'][idx].lerp_(
+                        outer_product.to(state['GG'][idx].dtype), weight=1.0 - state['shampoo_beta']
+                    )
+
+        if state['Q'] is None:
+            state['Q'] = self.get_orthogonal_matrix(state['GG'])
+
+        if step > 0 and step % state['precondition_frequency'] == 0:
+            state['Q'] = self.get_orthogonal_matrix_qr(state, max_precondition_dim, merge_dims)
+
+    @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 = group['betas']
+            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(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                    self.init_pre_conditioner(
+                        grad,
+                        state,
+                        precondition_frequency=group['precondition_frequency'],
+                        shampoo_beta=group['shampoo_beta'] if group['shampoo_beta'] is not None else beta2,
+                        max_precondition_dim=group['max_precondition_dim'],
+                        precondition_1d=group['precondition_1d'],
+                        merge_dims=group['merge_dims'],
+                    )
+
+                    self.update_pre_conditioner(
+                        grad,
+                        state,
+                        step=group['step'],
+                        max_precondition_dim=group['max_precondition_dim'],
+                        precondition_1d=group['precondition_1d'],
+                        merge_dims=group['merge_dims'],
+                    )
+
+                    continue
+
+                grad_projected = self.project(
+                    grad, state, merge_dims=group['merge_dims'], max_precondition_dim=group['max_precondition_dim']
+                )
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).add_(grad_projected.square(), alpha=1.0 - beta2)
+
+                de_nom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                exp_avg_projected = self.project(
+                    exp_avg, state, merge_dims=group['merge_dims'], max_precondition_dim=group['max_precondition_dim']
+                )
+
+                step_size = group['lr']
+                if group['correct_bias']:
+                    bias_correction1: float = self.debias(beta1, group['step'])
+                    bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+
+                    step_size *= bias_correction2_sq / bias_correction1
+
+                norm_grad = self.project(
+                    exp_avg_projected / de_nom,
+                    state,
+                    merge_dims=group['merge_dims'],
+                    max_precondition_dim=group['max_precondition_dim'],
+                    project_type='backward',
+                )
+
+                p.add_(norm_grad, alpha=-step_size)
+
+                self.apply_weight_decay(
+                    p,
+                    grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=True,
+                    fixed_decay=False,
+                )
+
+                self.update_pre_conditioner(
+                    grad,
+                    state,
+                    step=group['step'],
+                    max_precondition_dim=group['max_precondition_dim'],
+                    merge_dims=group['merge_dims'],
+                    precondition_1d=group['precondition_1d'],
+                )
+
+        return loss
diff --git a/requirements-dev.txt b/requirements-dev.txt
index d78222be0..1cf9e2de7 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -5,7 +5,7 @@ click==8.1.7 ; python_version >= "3.8" and python_full_version < "4.0.0"
 colorama==0.4.6 ; python_version >= "3.8" and python_full_version < "4.0.0" and (sys_platform == "win32" or platform_system == "Windows")
 coverage[toml]==7.6.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 exceptiongroup==1.2.2 ; python_version >= "3.8" and python_version < "3.11"
-filelock==3.16.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
+filelock==3.16.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 fsspec==2024.9.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
 iniconfig==2.0.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
 isort==5.13.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
@@ -17,12 +17,12 @@ networkx==3.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 numpy==1.24.4 ; python_version >= "3.8" and python_full_version < "4.0.0"
 packaging==24.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 pathspec==0.12.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
-platformdirs==4.3.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
+platformdirs==4.3.6 ; python_version >= "3.8" and python_full_version < "4.0.0"
 pluggy==1.5.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
 pytest-cov==5.0.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
-pytest==8.3.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
-ruff==0.6.4 ; python_version >= "3.8" and python_full_version < "4.0.0"
-sympy==1.13.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
+pytest==8.3.3 ; python_version >= "3.8" and python_full_version < "4.0.0"
+ruff==0.6.6 ; python_version >= "3.8" and python_full_version < "4.0.0"
+sympy==1.13.3 ; python_version >= "3.8" and python_full_version < "4.0.0"
 tomli==2.0.1 ; python_version >= "3.8" and python_full_version <= "3.11.0a6"
 torch==2.4.1+cpu ; python_version >= "3.8" and python_full_version < "4.0.0"
 typing-extensions==4.12.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
diff --git a/requirements.txt b/requirements.txt
index b530c3cd0..02e4831c0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,12 +1,12 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 
-filelock==3.16.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
+filelock==3.16.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 fsspec==2024.9.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
 jinja2==3.1.4 ; python_version >= "3.8" and python_full_version < "4.0.0"
 markupsafe==2.1.5 ; python_version >= "3.8" and python_full_version < "4.0.0"
 mpmath==1.3.0 ; python_version >= "3.8" and python_full_version < "4.0.0"
 networkx==3.1 ; python_version >= "3.8" and python_full_version < "4.0.0"
 numpy==1.24.4 ; python_version >= "3.8" and python_full_version < "4.0.0"
-sympy==1.13.2 ; python_version >= "3.8" and python_full_version < "4.0.0"
+sympy==1.13.3 ; python_version >= "3.8" and python_full_version < "4.0.0"
 torch==2.4.1+cpu ; python_version >= "3.8" and python_full_version < "4.0.0"
 typing-extensions==4.12.2 ; python_version >= "3.8" and python_full_version < "4.0.0"


Test patch:
diff --git a/tests/constants.py b/tests/constants.py
index 583df16ad..1a2d32803 100644
--- a/tests/constants.py
+++ b/tests/constants.py
@@ -13,6 +13,7 @@
     SGDP,
     SGDW,
     SM3,
+    SOAP,
     SRMM,
     SWATS,
     A2Grad,
@@ -140,6 +141,7 @@
     'adammini',
     'adamg',
     'ademamix',
+    'soap',
 ]
 
 VALID_LR_SCHEDULER_NAMES: List[str] = [
@@ -475,6 +477,11 @@
     (AdamG, {'lr': 1e0}, 20),
     (AdEMAMix, {'lr': 1e0}, 5),
     (AdEMAMix, {'lr': 1e0, 't_alpha_beta3': 5}, 5),
+    (
+        SOAP,
+        {'lr': 1e0, 'shampoo_beta': 0.95, 'precondition_frequency': 1, 'merge_dims': False, 'precondition_1d': True},
+        5,
+    ),
 ]
 ADANORM_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adanorm': True}, 10),
diff --git a/tests/test_gradients.py b/tests/test_gradients.py
index 14dd8be68..435e20d08 100644
--- a/tests/test_gradients.py
+++ b/tests/test_gradients.py
@@ -115,7 +115,7 @@ def test_sparse_supported(sparse_optimizer):
 
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)
 def test_bf16_gradient(optimizer_name):
-    if optimizer_name in {'shampoo', 'lomo', 'adalomo', 'bsam', 'adammini'}:
+    if optimizer_name in {'shampoo', 'lomo', 'adalomo', 'bsam', 'adammini', 'soap'}:
         pytest.skip(f'skip {optimizer_name}')
 
     param = torch.randn(1, 1).bfloat16().requires_grad_(True)
diff --git a/tests/test_load_modules.py b/tests/test_load_modules.py
index bf1d1c7a2..7fef074d1 100644
--- a/tests/test_load_modules.py
+++ b/tests/test_load_modules.py
@@ -38,7 +38,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 75
+    assert len(get_supported_optimizers()) == 76
 
 
 def test_get_supported_lr_schedulers():
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
index b862661ff..de9218dc6 100644
--- a/tests/test_optimizers.py
+++ b/tests/test_optimizers.py
@@ -92,7 +92,7 @@ def _closure() -> float:
 def test_lookahead(pullback_momentum, environment):
     (x_data, y_data), model, loss_fn = environment
 
-    optimizer = Lookahead(load_optimizer('adamp')(model.parameters(), lr=5e-1), pullback_momentum=pullback_momentum)
+    optimizer = Lookahead(load_optimizer('adamw')(model.parameters(), lr=5e-1), pullback_momentum=pullback_momentum)
 
     init_loss, loss = np.inf, np.inf
     for _ in range(5):
@@ -701,3 +701,50 @@ def test_trac_optimizer_erf_imag():
     optimizer.erf_imag(complex_tensor)
 
     assert str(optimizer).lower() == 'trac'
+
+
+@pytest.mark.parametrize(
+    'params',
+    [
+        {'merge_dims': True, 'precondition_1d': True, 'max_precondition_dim': 4, 'precondition_frequency': 1},
+        {'merge_dims': True, 'precondition_1d': False, 'max_precondition_dim': 1, 'precondition_frequency': 1},
+    ],
+)
+def test_soap_parameters(params, environment):
+    (x_data, y_data), _, loss_fn = environment
+
+    model = nn.Sequential(
+        nn.Linear(2, 8),
+        nn.Linear(8, 1),
+    )
+
+    optimizer = load_optimizer('soap')(model.parameters(), **params)
+
+    for _ in range(2):
+        optimizer.zero_grad()
+        loss_fn(model(x_data), y_data).backward()
+        optimizer.step()
+
+
+def test_soap_merge_dims_channel_last(environment):
+    (x_data, y_data), _, loss_fn = environment
+
+    x_data = x_data.reshape(-1, 1, 2, 1).repeat_interleave(2, dim=-1).to(memory_format=torch.channels_last)
+
+    model = nn.Sequential(
+        nn.Conv2d(1, 1, 2, 1),
+    )
+
+    optimizer = load_optimizer('soap')(
+        model.parameters(),
+        merge_dims=True,
+        precondition_1d=True,
+        max_precondition_dim=2,
+        precondition_frequency=1,
+        data_format='channels_last',
+    )
+
+    for _ in range(2):
+        optimizer.zero_grad()
+        loss_fn(model(x_data).squeeze(), y_data.squeeze()).backward()
+        optimizer.step()