heavyball 1.6.1__tar.gz → 1.6.3__tar.gz

{heavyball-1.6.1 → heavyball-1.6.3}/PKG-INFO

@@ -1,26 +1,32 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: heavyball
-Version: 1.6.1
-Summary: Efficient optimizers
-Home-page: https://github.com/HomebrewML/HeavyBall
-Author: HeavyBall Authors
-Author-email: github.heavyball@nestler.sh
-License: BSD
-Classifier: Development Status :: 5 - Production/Stable
-Classifier: License :: OSI Approved :: BSD License
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Topic :: Software Development :: Libraries
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Version: 1.6.3
+Summary: Efficient Optimizers
+Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
+Project-URL: source, https://github.com/HomebrewML/HeavyBall
+Project-URL: tracker, https://github.com/HomebrewML/HeavyBall/issues
+Keywords: torch,optimizer,muon,soap,psgd
 Classifier: Intended Audience :: Developers
-Requires-Python: >=3.7
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: opt-einsum
-Requires-Dist: torch
+Requires-Dist: opt-einsum>=3.4.0
+Requires-Dist: torch>=2.1.0
 Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: seaborn; extra == "dev"
+Requires-Dist: hyperopt; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: typer; extra == "dev"
 
 # `heavyball`: Efficient Optimizers
 

{heavyball-1.6.1 → heavyball-1.6.3}/heavyball/utils.py

@@ -376,7 +376,7 @@ def _compilable_scatter_set(target, source, index):
     target[:] = source.contiguous()[index].reshape_as(target)
 
 
-#@decorator_knowngood
+# @decorator_knowngood
 def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optional[Tensor] = None):
     """
     Computes the eigenbases of the preconditioner using one round of power iteration
@@ -426,7 +426,8 @@ def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optiona
     out_str = ''.join([o if o in to_shampoo else i for i, o in zip(in_str, out_str)])
 
     subscripts = f'{in_str},{from_shampoo},{to_shampoo}->{out_str}'
-    exp_avg_new = torch.einsum(subscripts, exp_avg, *[q for q in Q if q is not None], *[q for q in new_qs if q is not None])
+    exp_avg_new = torch.einsum(subscripts, exp_avg, *[q for q in Q if q is not None],
+                               *[q for q in new_qs if q is not None])
     copy_stochastic_(exp_avg, exp_avg_new)
 
     for q, q_new in zip(Q, new_qs):
@@ -434,7 +435,7 @@ def get_orthogonal_matrix_QR(GG: List[Tensor], Q: List[Tensor], exp_avg: Optiona
             copy_stochastic_(q, q_new)
 
 
-def get_orthogonal_matrix(mat):
+def get_orthogonal_matrix(mat, max_eps: float = 1e-3, min_eps: float = 1e-30):
     """
     Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
     """
@@ -448,23 +449,29 @@ def get_orthogonal_matrix(mat):
         m = promote(m.data)
 
         device, dtype = m.device, m.dtype
-        for modifier in (None, torch.double, 'cpu'):
-            if modifier is not None:
-                m = m.to(modifier)
+        eps = min_eps
+        while True:
             try:
-                eigval, eigvec = torch.linalg.eigh(m + 1e-30 * torch.eye(m.shape[0], device=m.device, dtype=m.dtype))
+                eye = torch.eye(m.shape[0], device=m.device, dtype=m.dtype)
+                eigval, eigvec = torch.linalg.eigh(m + eps * eye)
                 eigvec = eigvec.to(device=device, dtype=dtype)
                 break
             except torch.OutOfMemoryError:
-                pass
+                if m.device.type == 'cpu':
+                    raise
+                else:
+                    m = m.cpu()
             except RuntimeError:  # failed to compute eigenvalues
-                continue
+                if m.dtype != torch.double:
+                    m = m.double()
+                elif eps < max_eps:
+                    eps = eps ** (2 / 3)
+                else:
+                    raise
             clean()
-        else:
-            raise RuntimeError("Failed to compute eigenvalues.")
 
+        eigvec = eigvec.to(device=m.device, dtype=m.dtype)
         eigvec = torch.flip(eigvec, [1])
-
         final.append(eigvec)
 
     return final

{heavyball-1.6.1 → heavyball-1.6.3}/heavyball.egg-info/PKG-INFO

@@ -1,26 +1,32 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: heavyball
-Version: 1.6.1
-Summary: Efficient optimizers
-Home-page: https://github.com/HomebrewML/HeavyBall
-Author: HeavyBall Authors
-Author-email: github.heavyball@nestler.sh
-License: BSD
-Classifier: Development Status :: 5 - Production/Stable
-Classifier: License :: OSI Approved :: BSD License
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Topic :: Software Development :: Libraries
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Version: 1.6.3
+Summary: Efficient Optimizers
+Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
+Project-URL: source, https://github.com/HomebrewML/HeavyBall
+Project-URL: tracker, https://github.com/HomebrewML/HeavyBall/issues
+Keywords: torch,optimizer,muon,soap,psgd
 Classifier: Intended Audience :: Developers
-Requires-Python: >=3.7
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: opt-einsum
-Requires-Dist: torch
+Requires-Dist: opt-einsum>=3.4.0
+Requires-Dist: torch>=2.1.0
 Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: seaborn; extra == "dev"
+Requires-Dist: hyperopt; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: typer; extra == "dev"
 
 # `heavyball`: Efficient Optimizers
 

{heavyball-1.6.1 → heavyball-1.6.3}/heavyball.egg-info/SOURCES.txt

@@ -1,6 +1,6 @@
 LICENSE
 README.md
-setup.py
+pyproject.toml
 heavyball/__init__.py
 heavyball/chainable.py
 heavyball/utils.py

heavyball-1.6.3/heavyball.egg-info/requires.txt

@@ -0,0 +1,13 @@
+opt-einsum>=3.4.0
+torch>=2.1.0
+numpy
+
+[dev]
+pre-commit
+pytest
+ruff
+matplotlib
+seaborn
+hyperopt
+pandas
+typer

heavyball-1.6.3/pyproject.toml

@@ -0,0 +1,52 @@
+[build-system]
+requires = ["setuptools>=75.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "heavyball"
+description = "Efficient Optimizers"
+version = "1.6.3"
+authors = [{ name = "HeavyBall Authors", email = "github.heavyball@nestler.sh" }]
+classifiers = ["Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: BSD License",
+    "Natural Language :: English",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+]
+dependencies = ["opt-einsum>=3.4.0",
+    "torch>=2.1.0",
+    "numpy",
+]
+keywords = ["torch",
+    "optimizer",
+    "muon",
+    "soap",
+    "psgd",
+]
+readme = "README.md"
+requires-python = ">=3.9"
+
+[project.optional-dependencies]
+dev = ["pre-commit", "pytest", "ruff", "matplotlib", "seaborn", "hyperopt", "pandas", "typer"]
+
+[project.urls]
+source = "https://github.com/HomebrewML/HeavyBall"
+tracker = "https://github.com/HomebrewML/HeavyBall/issues"
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+extend-select = ["I", "W"]
+ignore = ["E741"]
+preview = true
+
+[tool.ruff.lint.isort]
+relative-imports-order = "closest-to-furthest"
+
+[tool.ruff.format]
+preview = true
+
+[tool.setuptools.packages.find]
+include = ["heavyball*"]

heavyball-1.6.1/heavyball.egg-info/requires.txt

@@ -1,3 +0,0 @@
-opt-einsum
-torch
-numpy

heavyball-1.6.1/setup.py

@@ -1,33 +0,0 @@
-import setuptools
-
-
-with open('README.md') as f:
-    README = f.read()
-
-setuptools.setup(
-    author="HeavyBall Authors",
-    author_email="github.heavyball@nestler.sh",
-    name='heavyball',
-    license='BSD',
-    description='Efficient optimizers',
-    version='1.6.1',
-    long_description=README,
-    url='https://github.com/HomebrewML/HeavyBall',
-    packages=setuptools.find_packages(),
-    python_requires=">=3.7",
-    long_description_content_type="text/markdown",
-    install_requires=['opt-einsum', 'torch', 'numpy'],
-    classifiers=[
-        # Trove classifiers
-        # (https://pypi.python.org/pypi?%3Aaction=list_classifiers)
-        'Development Status :: 5 - Production/Stable',
-        'License :: OSI Approved :: BSD License',
-        'Programming Language :: Python',
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
-        'Programming Language :: Python :: 3.9',
-        'Topic :: Software Development :: Libraries',
-        'Topic :: Software Development :: Libraries :: Python Modules',
-        'Intended Audience :: Developers',
-    ],
-)