<p align="center">
<img src="https://user-images.githubusercontent.com/6753639/251633368-1ec42732-1759-45d7-b949-51df6429a90a.svg" width="300" height="150">
</p>
<!--
<p align="center">
<img src="https://github.com/wilson-labs/cola/assets/6753639/28630ef8-5dcb-41c2-9f36-3cbba52f3d88.svg" width="300" height="139.29">
</p> -->
<!--
<p align = "center">
<img src="https://github.com/wilson-labs/cola/assets/6753639/8b02c51e-0e1e-44f5-a52a-47ad428688e4.svg" width="300" height="139.29">
</p>-->
# Compositional Linear Algebra (CoLA)
[](https://cola.readthedocs.io/en/stable/)
[](https://github.com/wilson-labs/cola/actions/workflows/python-package.yml)
[](https://codecov.io/gh/wilson-labs/cola)
[](https://pypi.org/project/cola-ml/)
[](https://arxiv.org/abs/2309.03060)
[](https://pepy.tech/project/cola-ml)
<!-- [](https://colab.research.google.com/github/wilson-labs/cola/blob/master/docs/notebooks/colabs/all.ipynb) -->
CoLA is a framework for scalable linear algebra, automatically exploiting the structure often found in machine learning problems and beyond.
CoLA natively supports PyTorch, JAX, as well as (limited) NumPy if JAX is not installed.
## Installation
```shell
pip install cola-ml
```
## Features in CoLA
* Large scale linear algebra routines for `solve(A,b)`, `eig(A)`, `logdet(A)`, `exp(A)`, `trace(A)`, `diag(A)`, `sqrt(A)`.
* Provides (user extendible) compositional rules to exploit structure through multiple dispatch.
* Has memory-efficient autodiff rules for iterative algorithms.
* Works with PyTorch or JAX, supporting GPU hardware acceleration.
* Supports operators with complex numbers and low precision.
* Provides linear algebra operations for both symmetric and non-symmetric matrices.
See https://cola.readthedocs.io/en/latest/ for our full documentation and many examples.
## Quick start guide
1. **LinearOperators**. The core object in CoLA is the LinearOperator. You can add and subtract them `+, -`,
multiply by constants `*, /`, matrix multiply them `@` and combine them in other ways:
`kron, kronsum, block_diag` etc.
```python
import jax.numpy as jnp
import cola
A = cola.ops.Diagonal(jnp.arange(5) + .1)
B = cola.ops.Dense(jnp.array([[2., 1.], [-2., 1.1], [.01, .2]]))
C = B.T @ B
D = C + 0.01 * cola.ops.I_like(C)
E = cola.ops.Kronecker(A, cola.ops.Dense(jnp.ones((2, 2))))
F = cola.ops.BlockDiag(E, D)
v = jnp.ones(F.shape[-1])
print(F @ v)
```
```
[0.2 0.2 2.2 2.2 4.2 4.2 6.2
6.2 8.2 8.2 7.8 2.1 ]
```
2. **Performing Linear Algebra**. With these objects we can perform linear algebra operations even when they are very big.
```python
print(cola.linalg.trace(F))
Q = F.T @ F + 1e-3 * cola.ops.I_like(F)
b = cola.linalg.inv(Q) @ v
print(jnp.linalg.norm(Q @ b - v))
print(cola.linalg.eig(F, k=F.shape[0])[0][:5])
print(cola.linalg.sqrt(A))
```
```
31.2701
0.0010193728
[ 2.0000000e-01+0.j 0.0000000e+00+0.j 2.1999998e+00+0.j
-1.1920929e-07+0.j 4.1999998e+00+0.j]
diag([0.31622776 1.0488088 1.4491377 1.7606816 2.0248456 ])
```
For many of these functions, if we know additional information about the matrices we can annotate them
to enable the algorithms to run faster.
```python
Qs = cola.SelfAdjoint(Q)
%timeit cola.linalg.inv(Q) @ v
%timeit cola.linalg.inv(Qs) @ v
```
3. **JAX and PyTorch**. We support both ML frameworks.
```python
import torch
A = cola.ops.Dense(torch.Tensor([[1., 2.], [3., 4.]]))
print(cola.linalg.trace(cola.kron(A, A)))
import jax.numpy as jnp
A = cola.ops.Dense(jnp.array([[1., 2.], [3., 4.]]))
print(cola.linalg.trace(cola.kron(A, A)))
```
```
tensor(25.)
25.0
```
CoLA also supports autograd (and jit):
```python
from jax import grad, jit, vmap
def myloss(x):
A = cola.ops.Dense(jnp.array([[1., 2.], [3., x]]))
return jnp.ones(2) @ cola.linalg.inv(A) @ jnp.ones(2)
g = jit(vmap(grad(myloss)))(jnp.array([.5, 10.]))
print(g)
```
```
[-0.06611571 -0.12499995]
```
## Citing us
If you use CoLA, please cite the following paper:
[Andres Potapczynski, Marc Finzi, Geoff Pleiss, and Andrew Gordon Wilson. "CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra." 2023.](https://arxiv.org/abs/2309.03060)
```
@article{potapczynski2023cola,
title={{CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra}},
author={Andres Potapczynski and Marc Finzi and Geoff Pleiss and Andrew Gordon Wilson},
journal={arXiv preprint arXiv:2309.03060},
year={2023}
}
```
### Features implemented
| Linear Algebra | inverse | eig | diag | trace | logdet | exp | sqrt | f(A) | SVD | pseudoinverse |
|:-----------------:|:-------:|:---:|:----:|:-----:|:------:|:---:|:----:|:--------:|:---:|:-------------:|
| **Implementation**| ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | |
| LinearOperators | Diag | BlockDiag | Kronecker | KronSum | Sparse | Jacobian | Hessian | Fisher | Concatenated | Triangular | FFT | Tridiagonal |
|:-----------------:|:----:|:---------:|:---------:|:-------:|:------:|:--------:|:-------:|:------:|:------------:|:----------:|:---:|:-----------:|
| **Implementation**| ✓ | ✓ | ✓ | ✓ | | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
| Annotations | SelfAdjoint | PSD | Unitary |
|:----------------:|:-----------:|:---:|:-------:|
| **Implementation**| ✓ | ✓ | ✓ |
| Backends | PyTorch | JAX | NumPy |
|:----------------:|:-----------:|:---:|:-------:|
| **Implementation**| ✓ | ✓ |Most operations|
## Contributing
See the contributing guidelines [docs/CONTRIBUTING.md](https://cola.readthedocs.io/en/latest/contributing.html) for information on submitting issues
and pull requests.
CoLA is Apache 2.0 licensed.
## Support and contact
Please raise an issue if you find a bug or slow performance when using CoLA.
Raw data
{
"_id": null,
"home_page": "https://github.com/wilson-labs/cola",
"name": "cola-ml",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.10",
"maintainer_email": null,
"keywords": "linear algebra, linear ops, sparse, PDE, AI, ML",
"author": "Marc Finzi and Andres Potapczynski",
"author_email": "maf820@nyu.edu",
"download_url": "https://files.pythonhosted.org/packages/67/fd/561b46c6f8f0837cb9fcf8af131e0ad2895d8b9bdf47b778a0ecd6f4a4ae/cola_ml-0.0.6.tar.gz",
"platform": null,
"description": "<p align=\"center\">\n <img src=\"https://user-images.githubusercontent.com/6753639/251633368-1ec42732-1759-45d7-b949-51df6429a90a.svg\" width=\"300\" height=\"150\">\n</p>\n\n<!--\n<p align=\"center\">\n <img src=\"https://github.com/wilson-labs/cola/assets/6753639/28630ef8-5dcb-41c2-9f36-3cbba52f3d88.svg\" width=\"300\" height=\"139.29\">\n</p> -->\n<!--\n<p align = \"center\">\n <img src=\"https://github.com/wilson-labs/cola/assets/6753639/8b02c51e-0e1e-44f5-a52a-47ad428688e4.svg\" width=\"300\" height=\"139.29\">\n</p>-->\n\n\n# Compositional Linear Algebra (CoLA)\n\n[](https://cola.readthedocs.io/en/stable/)\n[](https://github.com/wilson-labs/cola/actions/workflows/python-package.yml)\n[](https://codecov.io/gh/wilson-labs/cola)\n[](https://pypi.org/project/cola-ml/)\n[](https://arxiv.org/abs/2309.03060)\n[](https://pepy.tech/project/cola-ml)\n<!-- [](https://colab.research.google.com/github/wilson-labs/cola/blob/master/docs/notebooks/colabs/all.ipynb) -->\n\nCoLA is a framework for scalable linear algebra, automatically exploiting the structure often found in machine learning problems and beyond.\nCoLA natively supports PyTorch, JAX, as well as (limited) NumPy if JAX is not installed.\n\n## Installation\n```shell\npip install cola-ml\n```\n\n## Features in CoLA\n* Large scale linear algebra routines for `solve(A,b)`, `eig(A)`, `logdet(A)`, `exp(A)`, `trace(A)`, `diag(A)`, `sqrt(A)`.\n* Provides (user extendible) compositional rules to exploit structure through multiple dispatch.\n* Has memory-efficient autodiff rules for iterative algorithms.\n* Works with PyTorch or JAX, supporting GPU hardware acceleration.\n* Supports operators with complex numbers and low precision.\n* Provides linear algebra operations for both symmetric and non-symmetric matrices.\n\nSee https://cola.readthedocs.io/en/latest/ for our full documentation and many examples.\n\n\n## Quick start guide\n1. **LinearOperators**. The core object in CoLA is the LinearOperator. You can add and subtract them `+, -`,\nmultiply by constants `*, /`, matrix multiply them `@` and combine them in other ways:\n`kron, kronsum, block_diag` etc.\n```python\nimport jax.numpy as jnp\nimport cola\n\nA = cola.ops.Diagonal(jnp.arange(5) + .1)\nB = cola.ops.Dense(jnp.array([[2., 1.], [-2., 1.1], [.01, .2]]))\nC = B.T @ B\nD = C + 0.01 * cola.ops.I_like(C)\nE = cola.ops.Kronecker(A, cola.ops.Dense(jnp.ones((2, 2))))\nF = cola.ops.BlockDiag(E, D)\n\nv = jnp.ones(F.shape[-1])\nprint(F @ v)\n```\n```\n[0.2 0.2 2.2 2.2 4.2 4.2 6.2\n 6.2 8.2 8.2 7.8 2.1 ]\n```\n\n2. **Performing Linear Algebra**. With these objects we can perform linear algebra operations even when they are very big.\n```python\nprint(cola.linalg.trace(F))\nQ = F.T @ F + 1e-3 * cola.ops.I_like(F)\nb = cola.linalg.inv(Q) @ v\nprint(jnp.linalg.norm(Q @ b - v))\nprint(cola.linalg.eig(F, k=F.shape[0])[0][:5])\nprint(cola.linalg.sqrt(A))\n```\n\n```\n31.2701\n0.0010193728\n[ 2.0000000e-01+0.j 0.0000000e+00+0.j 2.1999998e+00+0.j\n -1.1920929e-07+0.j 4.1999998e+00+0.j]\ndiag([0.31622776 1.0488088 1.4491377 1.7606816 2.0248456 ])\n```\n\nFor many of these functions, if we know additional information about the matrices we can annotate them\nto enable the algorithms to run faster.\n\n```python\nQs = cola.SelfAdjoint(Q)\n%timeit cola.linalg.inv(Q) @ v\n%timeit cola.linalg.inv(Qs) @ v\n```\n\n3. **JAX and PyTorch**. We support both ML frameworks.\n```python\nimport torch\nA = cola.ops.Dense(torch.Tensor([[1., 2.], [3., 4.]]))\nprint(cola.linalg.trace(cola.kron(A, A)))\n\nimport jax.numpy as jnp\nA = cola.ops.Dense(jnp.array([[1., 2.], [3., 4.]]))\nprint(cola.linalg.trace(cola.kron(A, A)))\n```\n\n```\ntensor(25.)\n25.0\n```\n\nCoLA also supports autograd (and jit):\n```python\nfrom jax import grad, jit, vmap\n\n\ndef myloss(x):\n A = cola.ops.Dense(jnp.array([[1., 2.], [3., x]]))\n return jnp.ones(2) @ cola.linalg.inv(A) @ jnp.ones(2)\n\n\ng = jit(vmap(grad(myloss)))(jnp.array([.5, 10.]))\nprint(g)\n```\n\n```\n[-0.06611571 -0.12499995]\n```\n\n## Citing us\nIf you use CoLA, please cite the following paper:\n\n[Andres Potapczynski, Marc Finzi, Geoff Pleiss, and Andrew Gordon Wilson. \"CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra.\" 2023.](https://arxiv.org/abs/2309.03060)\n```\n@article{potapczynski2023cola,\n title={{CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra}},\n author={Andres Potapczynski and Marc Finzi and Geoff Pleiss and Andrew Gordon Wilson},\n journal={arXiv preprint arXiv:2309.03060},\n year={2023}\n}\n```\n\n### Features implemented\n\n| Linear Algebra | inverse | eig | diag | trace | logdet | exp | sqrt | f(A) | SVD | pseudoinverse |\n|:-----------------:|:-------:|:---:|:----:|:-----:|:------:|:---:|:----:|:--------:|:---:|:-------------:|\n| **Implementation**| \u2713 | \u2713 | \u2713 | \u2713 | \u2713 | \u2713 | \u2713 | \u2713 | | |\n\n| LinearOperators | Diag | BlockDiag | Kronecker | KronSum | Sparse | Jacobian | Hessian | Fisher | Concatenated | Triangular | FFT | Tridiagonal |\n|:-----------------:|:----:|:---------:|:---------:|:-------:|:------:|:--------:|:-------:|:------:|:------------:|:----------:|:---:|:-----------:|\n| **Implementation**| \u2713 | \u2713 | \u2713 | \u2713 | | \u2713 | \u2713 | \u2713 | \u2713 | \u2713 | \u2713 | \u2713 |\n\n| Annotations | SelfAdjoint | PSD | Unitary |\n|:----------------:|:-----------:|:---:|:-------:|\n| **Implementation**| \u2713 | \u2713 | \u2713 |\n\n\n| Backends | PyTorch | JAX | NumPy |\n|:----------------:|:-----------:|:---:|:-------:|\n| **Implementation**| \u2713 | \u2713 |Most operations|\n\n## Contributing\nSee the contributing guidelines [docs/CONTRIBUTING.md](https://cola.readthedocs.io/en/latest/contributing.html) for information on submitting issues\nand pull requests.\n\nCoLA is Apache 2.0 licensed.\n\n## Support and contact\nPlease raise an issue if you find a bug or slow performance when using CoLA.\n",
"bugtrack_url": null,
"license": "MIT",
"summary": null,
"version": "0.0.6",
"project_urls": {
"Homepage": "https://github.com/wilson-labs/cola"
},
"split_keywords": [
"linear algebra",
" linear ops",
" sparse",
" pde",
" ai",
" ml"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "aeb9d943a37d4215290f5b64b89526c81688ace6b70c14f08cf454ff547a9249",
"md5": "57cdd195bddd84f536cd3c5912688cca",
"sha256": "0931d7734fe53ef9c6f9c739e5596de1a0d1b85bd7869ddef37129d8eeffad13"
},
"downloads": -1,
"filename": "cola_ml-0.0.6-py3-none-any.whl",
"has_sig": false,
"md5_digest": "57cdd195bddd84f536cd3c5912688cca",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.10",
"size": 69991,
"upload_time": "2024-09-25T01:36:18",
"upload_time_iso_8601": "2024-09-25T01:36:18.889041Z",
"url": "https://files.pythonhosted.org/packages/ae/b9/d943a37d4215290f5b64b89526c81688ace6b70c14f08cf454ff547a9249/cola_ml-0.0.6-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "67fd561b46c6f8f0837cb9fcf8af131e0ad2895d8b9bdf47b778a0ecd6f4a4ae",
"md5": "389f0dfb3bdac702b2ffa8cb34a95e09",
"sha256": "a1f47fb08dea8af0f342ad0872f589b7744177982a2d7f8cb8f2b443e2747efa"
},
"downloads": -1,
"filename": "cola_ml-0.0.6.tar.gz",
"has_sig": false,
"md5_digest": "389f0dfb3bdac702b2ffa8cb34a95e09",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.10",
"size": 16591447,
"upload_time": "2024-09-25T01:36:21",
"upload_time_iso_8601": "2024-09-25T01:36:21.353297Z",
"url": "https://files.pythonhosted.org/packages/67/fd/561b46c6f8f0837cb9fcf8af131e0ad2895d8b9bdf47b778a0ecd6f4a4ae/cola_ml-0.0.6.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-09-25 01:36:21",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "wilson-labs",
"github_project": "cola",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "cola-ml"
}
JSON
