Introducing ConGrad: Backend-Agnostic Conjugate Gradients in Python
When working at scale, it is often slow or impossible to form full-sized matrices. Instead, we make use of black-box functions to calculate matrix-vector products \(Ab\). There is thus a need for solvers which accept these functions instead of requiring matrix representations of our linear maps – and, in particular, for conjugate gradient solvers for the symmetric positive definite case.
Historically, PyTorch users have used torch_cg, which gives correct results and has powerful batching support. However, it is somewhat poorly optimized, PyTorch specific, and currently unmaintained. All three of these problems are solved by my new package, PyConGrad. Following the PyTorch convention, the PIP package for PyConGrad is simply called congrad. You can also install congrad[numpy], congrad[cupy], or congrad[torch] to enforce NumPy, CuPy, or PyTorch as dependencies.
PyConGrad assumes that whatever vector representation you are using will respect NumPy-like operator overloading for basic arithmetic operations such as addition and scalar multiplication. Beyond that, it requires a dot product, a norm, and a few other functions, which are provided by a backend class. Backends are provided out of the box for NumPy, CuPy, and PyTorch, or you can write your own. They can also be used to achieve more customizable batching behavior than PyConGrad’s default expectation of arbitrarily many batch dimensions followed by one vector dimension at the end.
See the GitHub repo for more information and documentation.