Skip to main content

a KLU solver for JAX

Project description

KLUJAX

version: 0.4.1

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

CPU & float64

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 float32/complex64 arrays will be cast to float64/complex128!

Usage

The klujax library provides a single function solve(Ai, Aj, Ax, b), which solves for x in the sparse linear system Ax=b, where A is explicitly given in COO-format (Ai, Aj, Ax).

NOTE: the sparse matrix represented by (Ai, Aj, Ax) needs to be coalesced! KLUJAX provides a coalesce function (which unfortunately is not jax-jittable).

Supported shapes (? suffix means optional):

  • Ai: (n_nz,)
  • Aj: (n_nz,)
  • Ax: (n_lhs?, n_nz)
  • b: (n_lhs?, n_col, n_rhs?)
  • A (represented by (Ai, Aj, Ax)): (n_lhs?, n_col, n_col)

KLUJAX will automatically select a sensible way to act on underdefined dimensions of Ax and b:

dim(Ax) dim(b) assumed shape(Ax) assumed shape(b)
1D 1D n_nz n_col
1D 2D n_nz n_col x n_rhs
1D 3D n_nz n_lhs x n_col x n_rhs
2D 1D n_lhs x n_nz n_col
2D 2D n_lhs x n_nz n_lhs x n_col
2D 3D n_lhs x n_nz n_lhs x n_col x n_rhs

Where the A is always acting on the n_col dimension of b. The n_lhs dim is a shared batch dimension between A and b.

Additional dimensions can be added with jax.vmap (alternatively any higher dimensional problem can be reduced to the one above by properly transposing and reshaping Ax and b).

NOTE: JAX now has an experimental sparse library (jax.experimental.sparse). Using this natively in KLUJAX is not yet supported (but converting from BCOO or COO to Ai, Aj, Ax is trivial).

Basic Example

Script:

import klujax
import jax.numpy as jnp

b = jnp.array([8, 45, -3, 3, 19])
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],
    ]
)
Ai, Aj = jnp.where(jnp.abs(A_dense) > 0)
Ax = A_dense[Ai, Aj]

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

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

Output:

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

Installation

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 (see Installing from Source)

Installing from Source

NOTE: Installing from source should only be necessary when developing the library. If you as the user experience an install from source please create an issue.

Before installing, clone the build dependencies:

git clone --depth 1 --branch v7.2.0 https://github.com/DrTimothyAldenDavis/SuiteSparse suitesparse
git clone --depth 1 --branch main https://github.com/openxla/xla xla
git clone --depth 1 --branch stable https://github.com/pybind/pybind11 pybind11

Linux

On linux, you'll need gcc and g++, then inside the repo:

pip install .

MacOs

On MacOS, you'll need clang, then inside the repo:

pip install .

Windows

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 .

License & Credits

© Floris Laporte 2022, LGPL-2.1

This library was partly based on:

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.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

klujax-0.4.1-cp313-cp313-win_amd64.whl (139.7 kB view details)

Uploaded CPython 3.13 Windows x86-64

klujax-0.4.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB view details)

Uploaded CPython 3.13 manylinux: glibc 2.17+ x86-64

klujax-0.4.1-cp313-cp313-macosx_11_0_arm64.whl (198.2 kB view details)

Uploaded CPython 3.13 macOS 11.0+ ARM64

klujax-0.4.1-cp312-cp312-win_amd64.whl (139.7 kB view details)

Uploaded CPython 3.12 Windows x86-64

klujax-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB view details)

Uploaded CPython 3.12 manylinux: glibc 2.17+ x86-64

klujax-0.4.1-cp312-cp312-macosx_11_0_arm64.whl (198.1 kB view details)

Uploaded CPython 3.12 macOS 11.0+ ARM64

klujax-0.4.1-cp311-cp311-win_amd64.whl (139.2 kB view details)

Uploaded CPython 3.11 Windows x86-64

klujax-0.4.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB view details)

Uploaded CPython 3.11 manylinux: glibc 2.17+ x86-64

klujax-0.4.1-cp311-cp311-macosx_11_0_arm64.whl (199.4 kB view details)

Uploaded CPython 3.11 macOS 11.0+ ARM64

klujax-0.4.1-cp310-cp310-win_amd64.whl (137.8 kB view details)

Uploaded CPython 3.10 Windows x86-64

klujax-0.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB view details)

Uploaded CPython 3.10 manylinux: glibc 2.17+ x86-64

klujax-0.4.1-cp310-cp310-macosx_11_0_arm64.whl (197.9 kB view details)

Uploaded CPython 3.10 macOS 11.0+ ARM64

File details

Details for the file klujax-0.4.1-cp313-cp313-win_amd64.whl.

File metadata

  • Download URL: klujax-0.4.1-cp313-cp313-win_amd64.whl
  • Upload date:
  • Size: 139.7 kB
  • Tags: CPython 3.13, Windows x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.1-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 29aef61fec1fb427cc6a16cf09b63c06ba676c590aa1ea76ff90c768c1037209
MD5 18d1bf88084a3c85614f0ad56b7b7f8a
BLAKE2b-256 52b732fcb0aa90ea97e977d81cdc21cb8ae01a80409c6a475f9646fcc6eb4b51

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 cd85152975282f9e85c38fd6cc16c4da23562bff166c2e8c75a9b473a47d9429
MD5 03458188950c255fa2d9868e329b1030
BLAKE2b-256 a924456468fce588a38d73da746e90c13d1b01a83993dfd2ea46b7dc9ea3b983

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp313-cp313-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp313-cp313-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 e1a446db0cf70e754bb99f157dca64523fb6d822a0db649af028fff5d15ac39c
MD5 f8dfdf1836ecb1af8caa06174344217d
BLAKE2b-256 b43a34d16b3b497744b47af2defa9cae4856382e753c286a28e081971b386bf3

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp312-cp312-win_amd64.whl.

File metadata

  • Download URL: klujax-0.4.1-cp312-cp312-win_amd64.whl
  • Upload date:
  • Size: 139.7 kB
  • Tags: CPython 3.12, Windows x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.1-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 feb52cf1d932c4af2ceda1d1c56bbfd59e1b5e12697d7f3c8386d424cc5590ba
MD5 f74ee3386abef1cb9e923e46057ae2e6
BLAKE2b-256 d8097a3510a1ff46c65e18aa83634a1ff160c78928ec12d3f8276495d3abfebb

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 6cf1dcca3c2d15f5ee62d7bbb755788321a91d5e0e000aeebf0edf1aa4167b74
MD5 5d31ba19e12b482616da9d480066e71b
BLAKE2b-256 57f53a911205fda8248dab7d96411dfda039cb960b601657aa5e8f2bd2e83328

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp312-cp312-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 2b1d2da1745e45711fa9aa4e1348eb1a2a4f22586cff0890052180c0c6a45c60
MD5 4a7c80a121efdb497ba77de139ed530d
BLAKE2b-256 96bfa98babad40da2b7ea36a7f7d277a4c93903ad4b0b4d04c2d73d16cab47bd

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp311-cp311-win_amd64.whl.

File metadata

  • Download URL: klujax-0.4.1-cp311-cp311-win_amd64.whl
  • Upload date:
  • Size: 139.2 kB
  • Tags: CPython 3.11, Windows x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.1-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 5021291ee866209adda01bca8ca417ecd789b0c8673746ee1d02498ee6c61d57
MD5 494e5dfca47300bb2af10e2e642c80c3
BLAKE2b-256 764b64a389bf85d7a972787dcf92e854bbae74492f3d7f93c3601cd5d297082d

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 83b43a6975d85c6fa866f00d45498b40c38ece5aaac670b1698cd6ffec5040e5
MD5 48a79d4a9ea873a4530639df255713e5
BLAKE2b-256 433e9062e9d3594d3b48422cb388aa7c3f28d610a793583e713a685439753af5

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp311-cp311-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 32b0dcacaca3609b032c082ff53e15812e174bb663e07fd2955a731d471ba87d
MD5 b6b55f839a9da6400611858b45ed97f3
BLAKE2b-256 a917c32c8bfc54dedc413bbb441d6a8ffc212aaf017d10b9f32cd1b6469fb699

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp310-cp310-win_amd64.whl.

File metadata

  • Download URL: klujax-0.4.1-cp310-cp310-win_amd64.whl
  • Upload date:
  • Size: 137.8 kB
  • Tags: CPython 3.10, Windows x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.1-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 95e99fcdf4a9706973bc04bbbee84283e27f70bebf3859ded621b81758c1df41
MD5 6f01039d1bb7594dd375ab0a64283d67
BLAKE2b-256 cc0ec6c6c717af6100091445dff97863210e767e33a89d59c13247fb2f086d68

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 703ec664340efa58e90cca5a048d71d0b080c1c3b6a76e2e2625ee29e8095e68
MD5 42a05102d13b4b935887233435e71146
BLAKE2b-256 d92047f47108e9307c9aa94584b6052cc6cd1d068c1dc66bdd48f1640896b9fd

See more details on using hashes here.

File details

Details for the file klujax-0.4.1-cp310-cp310-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for klujax-0.4.1-cp310-cp310-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 9a415b32c378dcaa23f2f45e13796a740285d61a0414f3542965121a84f2a9c2
MD5 091024038a6155898987e8ab43dea079
BLAKE2b-256 e9ff3acd544314315136508def1eb17f8809d0b7f8c29698bcb0e4045ace4bdb

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page