Skip to main content

a KLU solver for JAX

Project description

KLUJAX

version: 0.4.3

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.3-cp313-cp313-win_arm64.whl (117.3 kB view details)

Uploaded CPython 3.13 Windows ARM64

klujax-0.4.3-cp313-cp313-win_amd64.whl (139.9 kB view details)

Uploaded CPython 3.13 Windows x86-64

klujax-0.4.3-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.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (2.0 MB view details)

Uploaded CPython 3.13 manylinux: glibc 2.17+ ARM64

klujax-0.4.3-cp313-cp313-macosx_11_0_arm64.whl (198.4 kB view details)

Uploaded CPython 3.13 macOS 11.0+ ARM64

klujax-0.4.3-cp312-cp312-win_arm64.whl (117.3 kB view details)

Uploaded CPython 3.12 Windows ARM64

klujax-0.4.3-cp312-cp312-win_amd64.whl (139.8 kB view details)

Uploaded CPython 3.12 Windows x86-64

klujax-0.4.3-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.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (2.0 MB view details)

Uploaded CPython 3.12 manylinux: glibc 2.17+ ARM64

klujax-0.4.3-cp312-cp312-macosx_11_0_arm64.whl (198.3 kB view details)

Uploaded CPython 3.12 macOS 11.0+ ARM64

klujax-0.4.3-cp311-cp311-win_arm64.whl (117.4 kB view details)

Uploaded CPython 3.11 Windows ARM64

klujax-0.4.3-cp311-cp311-win_amd64.whl (139.3 kB view details)

Uploaded CPython 3.11 Windows x86-64

klujax-0.4.3-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.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (2.0 MB view details)

Uploaded CPython 3.11 manylinux: glibc 2.17+ ARM64

klujax-0.4.3-cp311-cp311-macosx_11_0_arm64.whl (199.6 kB view details)

Uploaded CPython 3.11 macOS 11.0+ ARM64

klujax-0.4.3-cp310-cp310-win_arm64.whl (116.5 kB view details)

Uploaded CPython 3.10 Windows ARM64

klujax-0.4.3-cp310-cp310-win_amd64.whl (138.0 kB view details)

Uploaded CPython 3.10 Windows x86-64

klujax-0.4.3-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.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (2.0 MB view details)

Uploaded CPython 3.10 manylinux: glibc 2.17+ ARM64

klujax-0.4.3-cp310-cp310-macosx_11_0_arm64.whl (198.1 kB view details)

Uploaded CPython 3.10 macOS 11.0+ ARM64

File details

Details for the file klujax-0.4.3-cp313-cp313-win_arm64.whl.

File metadata

  • Download URL: klujax-0.4.3-cp313-cp313-win_arm64.whl
  • Upload date:
  • Size: 117.3 kB
  • Tags: CPython 3.13, Windows ARM64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.3-cp313-cp313-win_arm64.whl
Algorithm Hash digest
SHA256 cfe42f846f582a198442d7c4d882c7ca6c4f3f6af7891e54628aa035c8c744cc
MD5 e1b28e1131d3cf3c4fffa42af561ef83
BLAKE2b-256 18544b2fa898557604bf9a9410789dd69c84855249d6efebcb250f82ba3238f2

See more details on using hashes here.

File details

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

File metadata

  • Download URL: klujax-0.4.3-cp313-cp313-win_amd64.whl
  • Upload date:
  • Size: 139.9 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.3-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 26d82f832315753c20183d3ef1c37c0c3b3dedc6f0b68dabed1a13fe19ee5d1b
MD5 8a1e3b82fbcedb07bc54c77d83eb2a25
BLAKE2b-256 b624ee5e0ba4b4ae88b0cc1ff3715b3771c393f8756bc1486462a2ab78a7c314

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 d5fd7440b986bbfce7b65412af27066e3e186daf05be14452e5fe9e52bb32116
MD5 8c0bf6f0b780a71b4908905917cdbb03
BLAKE2b-256 b83e6eb7a96fe46c31b0431c80d2df82b40ccb2295e912d1b4bd58af2a37e6ec

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for klujax-0.4.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 8baeeefab615123dad4e13c3919354bf1da3c540a69f4440a86b6725950e752f
MD5 91baf609ae61d98670adafe5089474bc
BLAKE2b-256 67071d75736ab5d6be1c352907fef32976eadc21b7c1b3f6ca160a56db7defd8

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp313-cp313-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 36fbf836b92044aad9630072d5d84e655f1445c52992941873732f3ff9c3a3be
MD5 0daa73e3c429d3618c9347e687d8b1a9
BLAKE2b-256 964c3bf0f92f9d2702e990cc6d6ebc5be5cfc6bd1b3f2b789acb3f4a8d719f56

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp312-cp312-win_arm64.whl.

File metadata

  • Download URL: klujax-0.4.3-cp312-cp312-win_arm64.whl
  • Upload date:
  • Size: 117.3 kB
  • Tags: CPython 3.12, Windows ARM64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.3-cp312-cp312-win_arm64.whl
Algorithm Hash digest
SHA256 77e3a9c950d25ba7bdb767895b448457848a7e3a81ae9992622c544d012fbfc0
MD5 ef61ff8219b713a3e42c4abdcdb429e2
BLAKE2b-256 ab369f0447720edde9660ea880613f8fa75e128dcaf7ab04f56a369d5624ce41

See more details on using hashes here.

File details

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

File metadata

  • Download URL: klujax-0.4.3-cp312-cp312-win_amd64.whl
  • Upload date:
  • Size: 139.8 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.3-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 c7ab718933402dc4a4a58c3e10e3713d2e69626c6f20e67c495ada363a9fe171
MD5 a7b6f711f9ad5052c6821605e805299d
BLAKE2b-256 be8ef98b9335bed8415fd49bbe1204be5c8815679cb19190cc0e717d148782ef

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 b5b7a9e26d38adeaea658cbe844d7ac60fc8c6716b141d11c2ec43f65a976e15
MD5 f22809b838e3626b3dc6e86ba242b2be
BLAKE2b-256 0a6be22bab7854b4e687cd3723f0f1299abb64385ca0bc7bcc45dbf4474051d0

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for klujax-0.4.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 2dd76ecfbe3776298b2c637839c207bcad7993a3e5c039208b8a697350759a15
MD5 fad08a037c0ddc824b97e7d017367ded
BLAKE2b-256 bd5ece738c91af5ac476fc78c611a8a2fa42d1784b9613c8761e175086b402c5

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 fdf21636f4ceba4ebdb8337db8115099a9c8e85f7821015f4a869d608387ac2c
MD5 17c14d2e8a21fbfb7b03aa496489af90
BLAKE2b-256 78e2d93203ec169b8318dc21131f634989b51c6364b61dc26355a59d1c82c43d

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp311-cp311-win_arm64.whl.

File metadata

  • Download URL: klujax-0.4.3-cp311-cp311-win_arm64.whl
  • Upload date:
  • Size: 117.4 kB
  • Tags: CPython 3.11, Windows ARM64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.3-cp311-cp311-win_arm64.whl
Algorithm Hash digest
SHA256 54f2003f1eb662a8c8b632d014325dfdfb550d6c586634397ad871d4a2bd59c9
MD5 ce5c6c1e3018b642f1b93f06f6d8de77
BLAKE2b-256 bae287a73171cba7b79152367b5552ff80cef0faf2b953e3b65fc6a9ce12b021

See more details on using hashes here.

File details

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

File metadata

  • Download URL: klujax-0.4.3-cp311-cp311-win_amd64.whl
  • Upload date:
  • Size: 139.3 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.3-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 8799b6052f0da01382509fcf66e2e14bd24e11f3b498cc6acf15cd7752dc1c9b
MD5 b7fac5bceb1731fc238b3ee6dac8b50a
BLAKE2b-256 2f6af6858a6889394fe8f25c2433b2ada964bf365abdc496183846f16d8e11d0

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 0fe85fc53eb12dfd235e64f9b61a6db03e03e24217ec8eb955dda7875d97b0bc
MD5 188ee7fda02f624360408a6c19b7a9cc
BLAKE2b-256 887c707b6896130a3888e06675e4a80692fa5ebd85d93bdc62e0d99a085c6bd1

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for klujax-0.4.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 1ec6177494088e6ee27bfead210f1ad67fa6e6291eb9a259cdddf01a924d30e6
MD5 584e66c56459e20492c3789e16f277db
BLAKE2b-256 8d31e2716c8145010a70027d766774802308a47324d67f5042c1398acc63a887

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 c24eefef349a8d570a04a05a861499b7a7ed87ae2a537dac7f1744ab70dd2668
MD5 9d00cbecdc54715c5f18318b2617a6ee
BLAKE2b-256 849f93a2858053a3cdc30eceb10d3a42bfc2307426d4a48f246afac65c01fd17

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp310-cp310-win_arm64.whl.

File metadata

  • Download URL: klujax-0.4.3-cp310-cp310-win_arm64.whl
  • Upload date:
  • Size: 116.5 kB
  • Tags: CPython 3.10, Windows ARM64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.9

File hashes

Hashes for klujax-0.4.3-cp310-cp310-win_arm64.whl
Algorithm Hash digest
SHA256 0d25e5c57ea6d5eeb9d786f705bc2ecf9db5853a0740dbda2aa47bac5d80af7c
MD5 e3550fae94a05c29fee209e78554b5c4
BLAKE2b-256 097c4047c7de47a70ac543c0fa1ffa0fa5126d38dd9d1b0fa2168139bd1db038

See more details on using hashes here.

File details

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

File metadata

  • Download URL: klujax-0.4.3-cp310-cp310-win_amd64.whl
  • Upload date:
  • Size: 138.0 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.3-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 ee565eb8817466e0ca3f29f94856769cd083ecd402f2f95e839e74dee15dda44
MD5 5b0206043074eb29ed0da6d6a5dc7853
BLAKE2b-256 717b5c1a425a2a721ad0e4c7e1cd3b85c974403f2963fd88bbfa6eb6e4aa60ca

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 93c820ae1d02fa8b5b23aba7242087b7557dbc6b1478689d515ce4ab793e28f8
MD5 ee65cefdca5c4a19f2c4f54eb659cba3
BLAKE2b-256 dbe616baa36dbf2fce97acdf104ee2d18f7d04618caaf00ec1fdc2e28c61a068

See more details on using hashes here.

File details

Details for the file klujax-0.4.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for klujax-0.4.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 d8e5cd54acdd452407343d52076aa54cf33538ef7a3c57a66643d5a9d90b1c75
MD5 d6fb6131de8a171f1fdd65e18a88c8d2
BLAKE2b-256 0d1fcd826ad3ce32284736a639b04b2b9b1a6456e5aeab040d3f70e2563ec4a2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for klujax-0.4.3-cp310-cp310-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 26127b2f76f76d76e26bd20ea516e3b42cab020b6eeafc7bdb17c1e9258b11e1
MD5 d026bc5a2ba7d49870e682c7e87b4e54
BLAKE2b-256 0cf2679aea4dcb4dcbaeef75c6a19bb3e90817ed5c7c39165eaf05dacfb4ee17

See more details on using hashes here.

Supported by

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