heavyball 1.7.2__py3-none-any.whl → 2.0.0.dev0__py3-none-any.whl

heavyball-2.0.0.dev0.dist-info/METADATA

@@ -0,0 +1,109 @@
+Metadata-Version: 2.4
+Name: heavyball
+Version: 2.0.0.dev0
+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
+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>=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"
+Dynamic: license-file
+
+# heavyball
+
+[![PyPI version](https://img.shields.io/pypi/v/heavyball?color=blue)][pypi]
+[![License](https://img.shields.io/badge/license-BSD--3--Clause-blue.svg)][license]
+
+_High-performance, extensible, chainable optimizers for PyTorch._
+
+## Why heavyball
+
+- **Lightning-Fast Training**: Batched `foreach` operations deliver significant speedups on large models.
+- **Adaptive & Extensible**: Built-in AdamW, RMSprop, Schedule-Free algorithms, and PaLM-inspired schedules.
+- **Plug-and-Play**: Drop-in replacements for `torch.optim` with seamless integration.
+- **Customizable**: Chainable API lets you compose optimizers and transforms (MARS correction, cautious updates, orthogonal updates).
+- **Battle-Tested**: Extensive benchmarks and real-world examples included.
+
+## Key Features
+
+- Foreach-based optimizers: `ForeachAdamW`, `ForeachRMSprop`, `ForeachSFAdamW`, `Muon`, `ADOPT`, `MSAM`, …
+- Schedule-Free optimizers with dynamic learning rate adaptation.
+- Advanced update rules: MARS correction, cautious updates, PaLM beta2 scheduling.
+- Chainable transforms for custom optimization recipes.
+- Comprehensive benchmark suite (`benchmark/`).
+- Detailed documentation and example-driven tutorials.
+
+## Quickstart
+
+**Install:**
+```bash
+pip install heavyball
+```
+
+**Basic usage:**
+```python
+import torch
+from torch import nn
+from heavyball import ForeachAdamW
+
+model = nn.Sequential(
+    nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10)
+)
+optimizer = ForeachAdamW(model.parameters(), lr=1e-3)
+
+for data, target in dataloader:
+    optimizer.zero_grad()
+    output = model(data)
+    loss = torch.nn.functional.cross_entropy(output, target)
+    loss.backward()
+    optimizer.step()
+```
+
+## Benchmarks
+
+> Reproduce benchmarks with:
+> ```bash
+> python3 -m benchmark.run_all_benchmarks --opt ForeachSOAP --opt LaProp --opt AdamW --opt Muon --opt ForeachCachedNewtonPSGD  --opt RMSprop --opt OrthoLaProp --opt ForeachSFAdamW --opt ForeachADOPT --opt LaPropOrtho --opt CachedPSGDKron --opt SignLaProp --opt ForeachSOLP --opt PSGDLRA --opt NewtonPSGDLRA --opt NewtonHybrid2PSGDKron --opt NewtonHybrid2PSGDLRA --opt mars-NewtonHybrid2PSGDLRA --opt MSAMLaProp --opt mars-adaptive-NewtonHybrid2PSGDKron  --opt mars-ortho-NewtonHybrid2PSGDKron --opt MuonLaProp --opt mars-unscaled-NewtonHybrid2PSGDKron --opt mars-NewtonHybrid2PSGDKron --opt cautious-AdamW --opt unscaled_cautious-AdamW --opt mars-AdamW  --dtype float32 --steps 1000000 --trials 1000 --parallelism 256 --seeds 1 --difficulties trivial --difficulties easy --difficulties medium --difficulties hard --difficulties extreme --difficulties nightmare --timeout 2880
+> ```
+
+
+## Contributing
+
+We welcome contributions! Please check the [issue tracker][tracker] and follow these steps:
+1. Fork the repo and create a feature branch.
+2. Install dev dependencies: `pip install -e .[dev]`.
+3. Run tests: `pytest`.
+4. Submit a pull request.
+
+## License
+
+BSD 3-Clause — see the [LICENSE](LICENSE) file.
+
+---
+<p align="center">
+  Made by the HeavyBall team.
+</p>
+
+[pypi]: https://pypi.org/project/heavyball/
+[license]: LICENSE
+[tracker]: https://github.com/HomebrewML/HeavyBall/issues

heavyball-2.0.0.dev0.dist-info/RECORD

@@ -0,0 +1,9 @@
+heavyball/__init__.py,sha256=BUKhEKbpIY-93nMd5CEkhKv6d6LbSAstY_WSpFGGzQg,26017
+heavyball/chainable.py,sha256=aoPG_MtVncDqZRKeKsIkmnTsYLhmF3_I_06ZMovgTKc,39190
+heavyball/helpers.py,sha256=gKUhzu38f8e7CdnBbR8M51g9w8w0Kft_RW3V2fMXKw0,30159
+heavyball/utils.py,sha256=jrFHeZ0MjkiUutuLTT8ky8G_kaDFxIZACsfK_tRVQaE,90204
+heavyball-2.0.0.dev0.dist-info/licenses/LICENSE,sha256=G9fFZcNIVWjU7o6Pr_4sJBRCNDU5X-zelSxIJ2D48ms,1323
+heavyball-2.0.0.dev0.dist-info/METADATA,sha256=Au0Wdlt4YpTMnVRZdH29Svjwbuqks4fxivdFRXWPP7I,4254
+heavyball-2.0.0.dev0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+heavyball-2.0.0.dev0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-2.0.0.dev0.dist-info/RECORD,,

{heavyball-1.7.2.dist-info → heavyball-2.0.0.dev0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.7.1)
 Root-Is-Purelib: true
 Tag: py3-none-any