pytorch-opt
Implementations
The library implements Muon, Shampoo, SOAP, K-FAC, EKFAC, natural gradient descent, trust-region Newton-CG, Hessian-free optimization, AdaHessian, Sophia, and PSGD. It includes shared routines for matrix roots, Fisher and Hessian-vector products, Kronecker factors, and damping. The reference implementations use PyTorch. An optional C++/ATen backend supports CPU and CUDA.
Comparisons and diagnostics
Parameter routing keeps the AdamW fallback consistent between optimizer runs. Diagnostics report conditioning, damping, update timing, and optimizer state. The verification notes cover numerical checks, convergence tests, checkpoint round trips, and known limitations. I'm using this library for my pretraining optimizer experiments.