├── .github └── workflows │ ├── ci.yaml │ └── release.yaml ├── .gitignore ├── LICENSE ├── README.md ├── assets ├── benchmark_matrix.png ├── delayed_psgd.png └── memory.png ├── docs ├── assets │ ├── benchmark_matrix.png │ ├── early_stopping.png │ ├── hiker_analogy_diagram.png │ ├── psgd_efficiency_cache.png │ ├── psgd_efficiency_cache_triu_as_line.png │ ├── psgd_efficiency_triu_as_line.png │ └── saddle_point_comparison.gif ├── benchmark.md ├── heavyball2.md └── psgd_efficiency.md ├── examples ├── autoencoder.py ├── branched_optimizer.py └── modify_functions.py ├── heavyball ├── __init__.py ├── chainable.py ├── helpers.py └── utils.py ├── interactive ├── playground.py └── static │ ├── init-globals.js │ └── node-editor.js ├── pre-commit.yaml ├── pyproject.toml ├── scripts ├── migrate_optimizer_state.py └── run_ci_tests.sh └── test ├── benchmark_precond_fitting.py ├── benchmark_psgd_lb.py ├── readme.md ├── test_ademamix.py ├── test_bf16_params.py ├── test_bf16_q.py ├── test_bf16_storage.py ├── test_caution.py ├── test_chainable_cpu.py ├── test_channels_last.py ├── test_clip.py ├── test_closure.py ├── test_cpu_features.py ├── test_ema.py ├── test_foreach.py ├── test_helpers_cpu.py ├── test_hook.py ├── test_mars.py ├── test_memory.py ├── test_memory_leak.py ├── test_merge.py ├── test_migrate_cli.py ├── test_nd_param.py ├── test_no_grad.py ├── test_optimizer_cpu_smoke.py ├── test_psgd_precond_init_stability.py ├── test_save_restore.py ├── test_singular_values.py ├── test_soap.py ├── test_stochastic_updates.py ├── test_stochastic_utils_cpu.py ├── test_toy_training.py ├── test_utils_cpu.py ├── test_utils_property.py └── utils.py /.github/workflows/ci.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/.github/workflows/ci.yaml -------------------------------------------------------------------------------- /.github/workflows/release.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/.github/workflows/release.yaml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/README.md -------------------------------------------------------------------------------- /assets/benchmark_matrix.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/assets/benchmark_matrix.png -------------------------------------------------------------------------------- /assets/delayed_psgd.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/assets/delayed_psgd.png -------------------------------------------------------------------------------- /assets/memory.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/assets/memory.png -------------------------------------------------------------------------------- /docs/assets/benchmark_matrix.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/benchmark_matrix.png -------------------------------------------------------------------------------- /docs/assets/early_stopping.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/early_stopping.png -------------------------------------------------------------------------------- /docs/assets/hiker_analogy_diagram.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/hiker_analogy_diagram.png -------------------------------------------------------------------------------- /docs/assets/psgd_efficiency_cache.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/psgd_efficiency_cache.png -------------------------------------------------------------------------------- /docs/assets/psgd_efficiency_cache_triu_as_line.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/psgd_efficiency_cache_triu_as_line.png -------------------------------------------------------------------------------- /docs/assets/psgd_efficiency_triu_as_line.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/psgd_efficiency_triu_as_line.png -------------------------------------------------------------------------------- /docs/assets/saddle_point_comparison.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/assets/saddle_point_comparison.gif -------------------------------------------------------------------------------- /docs/benchmark.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/benchmark.md -------------------------------------------------------------------------------- /docs/heavyball2.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/heavyball2.md -------------------------------------------------------------------------------- /docs/psgd_efficiency.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/docs/psgd_efficiency.md -------------------------------------------------------------------------------- /examples/autoencoder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/examples/autoencoder.py -------------------------------------------------------------------------------- /examples/branched_optimizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/examples/branched_optimizer.py -------------------------------------------------------------------------------- /examples/modify_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/examples/modify_functions.py -------------------------------------------------------------------------------- /heavyball/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/heavyball/__init__.py -------------------------------------------------------------------------------- /heavyball/chainable.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/heavyball/chainable.py -------------------------------------------------------------------------------- /heavyball/helpers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/heavyball/helpers.py -------------------------------------------------------------------------------- /heavyball/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/heavyball/utils.py -------------------------------------------------------------------------------- /interactive/playground.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/interactive/playground.py -------------------------------------------------------------------------------- /interactive/static/init-globals.js: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/interactive/static/init-globals.js -------------------------------------------------------------------------------- /interactive/static/node-editor.js: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/interactive/static/node-editor.js -------------------------------------------------------------------------------- /pre-commit.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/pre-commit.yaml -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/pyproject.toml -------------------------------------------------------------------------------- /scripts/migrate_optimizer_state.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/scripts/migrate_optimizer_state.py -------------------------------------------------------------------------------- /scripts/run_ci_tests.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/scripts/run_ci_tests.sh -------------------------------------------------------------------------------- /test/benchmark_precond_fitting.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/benchmark_precond_fitting.py -------------------------------------------------------------------------------- /test/benchmark_psgd_lb.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/benchmark_psgd_lb.py -------------------------------------------------------------------------------- /test/readme.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/readme.md -------------------------------------------------------------------------------- /test/test_ademamix.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_ademamix.py -------------------------------------------------------------------------------- /test/test_bf16_params.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_bf16_params.py -------------------------------------------------------------------------------- /test/test_bf16_q.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_bf16_q.py -------------------------------------------------------------------------------- /test/test_bf16_storage.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_bf16_storage.py -------------------------------------------------------------------------------- /test/test_caution.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_caution.py -------------------------------------------------------------------------------- /test/test_chainable_cpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_chainable_cpu.py -------------------------------------------------------------------------------- /test/test_channels_last.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_channels_last.py -------------------------------------------------------------------------------- /test/test_clip.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_clip.py -------------------------------------------------------------------------------- /test/test_closure.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_closure.py -------------------------------------------------------------------------------- /test/test_cpu_features.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_cpu_features.py -------------------------------------------------------------------------------- /test/test_ema.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_ema.py -------------------------------------------------------------------------------- /test/test_foreach.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_foreach.py -------------------------------------------------------------------------------- /test/test_helpers_cpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_helpers_cpu.py -------------------------------------------------------------------------------- /test/test_hook.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_hook.py -------------------------------------------------------------------------------- /test/test_mars.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_mars.py -------------------------------------------------------------------------------- /test/test_memory.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_memory.py -------------------------------------------------------------------------------- /test/test_memory_leak.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_memory_leak.py -------------------------------------------------------------------------------- /test/test_merge.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_merge.py -------------------------------------------------------------------------------- /test/test_migrate_cli.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_migrate_cli.py -------------------------------------------------------------------------------- /test/test_nd_param.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_nd_param.py -------------------------------------------------------------------------------- /test/test_no_grad.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_no_grad.py -------------------------------------------------------------------------------- /test/test_optimizer_cpu_smoke.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_optimizer_cpu_smoke.py -------------------------------------------------------------------------------- /test/test_psgd_precond_init_stability.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_psgd_precond_init_stability.py -------------------------------------------------------------------------------- /test/test_save_restore.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_save_restore.py -------------------------------------------------------------------------------- /test/test_singular_values.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_singular_values.py -------------------------------------------------------------------------------- /test/test_soap.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_soap.py -------------------------------------------------------------------------------- /test/test_stochastic_updates.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_stochastic_updates.py -------------------------------------------------------------------------------- /test/test_stochastic_utils_cpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_stochastic_utils_cpu.py -------------------------------------------------------------------------------- /test/test_toy_training.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_toy_training.py -------------------------------------------------------------------------------- /test/test_utils_cpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_utils_cpu.py -------------------------------------------------------------------------------- /test/test_utils_property.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/test_utils_property.py -------------------------------------------------------------------------------- /test/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/HeavyBall/HEAD/test/utils.py --------------------------------------------------------------------------------