A sparse linear solver for JAX based on the efficient KLU algorithm.

This library is a wrapper around the SuiteSparse KLU algorithms. This means the algorithm is only implemented for C-arrays and hence is only available for CPU arrays with double precision, i.e. float64 or complex128.

Note that this will be enforced at import of klujax!

The klujax library provides a single function solve(A, b), which solves for x in the linear system Ax=b A is a sparse tensor in COO-format with shape mxm and x and b have shape mxn. Note that JAX does not have a native sparse matrix representation and hence A should be represented as a tuple of two index arrays and a value array: (Ai, Aj, Ax).

import jax.numpy as jnp
from klujax import solve

b = jnp.array([8, 45, -3, 3, 19], dtype=jnp.float64)
A_dense = jnp.array([[2, 3, 0, 0, 0],
                     [3, 0, 4, 0, 6],
                     [0, -1, -3, 2, 0],
                     [0, 0, 1, 0, 0],
                     [0, 4, 2, 0, 1]], dtype=jnp.float64)
Ai, Aj = jnp.where(jnp.abs(A_dense) > 0)
Ax = A_dense[Ai, Aj]

result_ref = jnp.linalg.inv(A_dense)@b
result = solve(Ai, Aj, Ax, b)

print(jnp.abs(result - result_ref) < 1e-12)
print(result)

[ True True True True True]
[1. 2. 3. 4. 5.]

The library is statically linked to the SuiteSparse C++ library. It can be installed on most platforms as follows:

pip install klujax

There exist pre-built wheels for Linux and Windows (python 3.8+). If no compatible wheel is found, however, pip will attempt to install the library from source... make sure you have the necessary build dependencies installed.

On linux, you'll need gcc and g++. Then just do a normal pip install:

pip install klujax

On Windows, installing from source is a bit more involved as typically the build dependencies are not installed. To install those, download Visual Studio Community 2017 from here. During installation, go to Workloads and select the following workloads:

• Desktop development with C++
• Python development

Then go to Individual Components and select the following additional items:

• C++/CLI support
• VC++ 2015.3 v14.00 (v140) toolset for desktop

Then, download and install Microsoft Visual C++ Redistributable from here.

After these installation steps, run the following commands inside a x64 Native Tools Command Prompt for VS 2017:

set DISTUTILS_USE_SDK=1
pip install klujax

© Floris Laporte 2022, LGPL-2.1

This library was partly based on:

• torch_sparse_solve, LGPL-2.1
• SuiteSparse, LGPL-2.1
• kagami-c/PyKLU, LGPL-2.1
• scipy.sparse, BSD-3

This library vendors an unmodified version of the SuiteSparse libraries in its source (.tar.gz) distribution to allow for static linking. This is in accordance with their LGPL licence.