KLUJAX

version: 0.5.0.post3

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

This is a fork of the original klujax package, meant for use in splineax. The aim is to eventually merge the changes here into the upstream package and remove this fork. In the meantime, the version number will stay as 0.5.0.postN, and N will be incremented sequentially on each new published version.

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!

Basic Usage

The klujax library provides a basic 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.]

Advanced Usage

For high-performance applications like transient simulations or iterative solvers, you should avoid using the high-level klujax.solve function. The klujax.solve is in fact a wrapper around three distinct parts of the KLU algorithm:

Analyze (Symbolic): Inspects the sparsity pattern ($A_i, A_j$) to find optimal permutations and block triangular forms. This depends only on the structure of the matrix.
Factorize (Numeric): Performs the actual LU decomposition. This depends on the values ($A_x$) and requires a symbolic handle.
Solve (Numeric): Executes forward and backward substitution to find $x$. This depends on the right-hand side ($b$) and requires a numeric handle.

Significant performance gains are achieved by hoisting the "Analysis" or "Factorization" steps out of your inner loops.

1. High-Performance Transient Pattern (Reusing Symbolic)

In a simulation where the sparsity pattern is constant but the values ($A_x$) and right-hand side ($b$) change, you should perform the expensive analyze step exactly once outside your JIT loop.

import jax
import klujax

# 1. Analyze once in Python (CPU)
# Returns a KLUHandleManager that automatically cleans up C++ memory
symbolic = klujax.analyze(Ai, Aj, n_col)

@jax.jit
def simulation_step(Ax_t, b_t, sym):
    # 2. Use the symbolic handle inside JIT
    # The solver will perform numeric factorization and solve
    return klujax.solve_with_symbol(Ai, Aj, Ax_t, b_t, sym)

for t in range(steps):
    x_t = simulation_step(Ax[t], b[t], symbolic)

Fine-Grained Control (Numeric Factorization)

If you need to solve the same system with many different $b$ vectors while the matrix $A$ remains constant, you can further split the numeric factorization. This is often performed in a modified Newton-Raphson loop where the computationally expensive jacobian+factorization is only evaluated once and the solve stage is deemed "cheap" in comparison

# Factorize the matrix once
numeric = klujax.factor(Ai, Aj, Ax, symbolic)

@jax.jit
def fast_solve(b_t, num, sym):
    # This call is extremely fast as it skips factorization entirely
    return klujax.solve_with_numeric(num, b_t, sym)

for i in range(100):
    x_i = fast_solve(b_batch[i], numeric, symbolic)

Safe Refactorization (Status Codes & Conditioning)

klujax.refactor reuses the pivot order picked for the original matrix. That is what makes it fast, but it also means the factorization can fail, or silently lose accuracy, once the values have drifted far enough. Two additions make that recoverable.

klujax.refactor_with_status reports a failure through a status code rather than raising, so it can be branched on under jax.jit where an error would abort everything. klujax.refactor_and_solve_with_status does the same for the fused path.

numeric, status = klujax.refactor_with_status(Ai, Aj, Ax_new, numeric, symbolic)
if status[0] != klujax.KLUStatus.OK:
    # the numeric object is unusable for a solve, but the symbolic one is fine
    klujax.free_numeric(numeric)
    numeric = klujax.factor(Ai, Aj, Ax_new, symbolic)

When status != OK the numeric object may be partially overwritten and must not be used for a solve. It remains valid to pass to free_numeric, and the symbolic object is unaffected and may be reused for a fresh factor.

Degradation short of outright failure is caught by klujax.rcond, the reciprocal pivot growth estimate min|Uii| / max|Uii|. It costs O(n), far less than solving with a probe vector and measuring the residual. klujax.condest gives a proper 1-norm condition number estimate and is the usual follow-up when rcond is borderline.

if klujax.rcond(symbolic, numeric)[0] < 1e-10:
    ...  # pivots have degraded, re-factor from scratch

Lifecycle & Pointer Pitfalls

Because klujax.analyze and klujax.factor generate KLUHandleManager objects which wrap low level C++ pointers, there are strict rules for avoiding memory leaks and segmentation faults.

The "Ghost Pointer" Problem inside JIT:

JAX's jit works by tracing your code. During tracing, Python objects like the KLUHandleManager are converted into symbolic Tracers.

Outside JIT: The KLUHandleManager uses RAII (Resource Acquisition Is Initialization). When the Python variable is deleted or goes out of scope, the C++ memory is freed automatically.
Inside JIT: If you create a handle (via analyze or factor) inside a JIT-compiled function, the Python manager is "lost" during the conversion to XLA. XLA will allocate the C++ memory at runtime, but it will never call the free function.

The Fix: Explicit Destruction with Dependencies

If you must create a handle inside JIT, you must manually call free_symbolic or free_numeric inside that same function. To prevent the compiler from freeing the pointer before the solve is finished, you must pass the solution as a dependency.

@jax.jit
def dynamic_solve(Ai, Aj, Ax, b):
    # 1. Born inside JIT (No automatic cleanup!)
    sym = klujax.analyze(Ai, Aj, 5)

    # 2. Compute solution
    x = klujax.solve_with_symbol(Ai, Aj, Ax, b, sym)

    # 3. CRITICAL: Force XLA to free 'sym' ONLY AFTER 'x' is ready
    klujax.free_symbolic(sym, dependency=x)

    return x

Summary of Best Practices

Hoist Creations: Always try to call analyze or factor outside of JIT blocks.
One Manager, One Free: Do not manually call free_symbolic(manager) and then let the manager go out of scope; it will attempt a double-free (though the library has safeguards to prevent a crash).
Check for Warnings: If you see a UserWarning: Allocating KLU handle inside JIT, your code is currently leaking memory. Use the dependency pattern shown above to fix it.

Installation

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

pip install splineax-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

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.

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

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

splineax_klujax-0.5.0.post3-cp314-cp314-win_arm64.whl (192.4 kB view details)

Uploaded Aug 12, 2026 CPython 3.14Windows ARM64

splineax_klujax-0.5.0.post3-cp314-cp314-win_amd64.whl (207.9 kB view details)

Uploaded Aug 12, 2026 CPython 3.14Windows x86-64

splineax_klujax-0.5.0.post3-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl (3.3 MB view details)

Uploaded Aug 12, 2026 CPython 3.14manylinux: glibc 2.26+ x86-64manylinux: glibc 2.28+ x86-64

splineax_klujax-0.5.0.post3-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (3.2 MB view details)

Uploaded Aug 12, 2026 CPython 3.14manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

splineax_klujax-0.5.0.post3-cp314-cp314-macosx_11_0_arm64.whl (277.4 kB view details)

Uploaded Aug 12, 2026 CPython 3.14macOS 11.0+ ARM64

splineax_klujax-0.5.0.post3-cp313-cp313-win_arm64.whl (186.2 kB view details)

Uploaded Aug 12, 2026 CPython 3.13Windows ARM64

splineax_klujax-0.5.0.post3-cp313-cp313-win_amd64.whl (202.9 kB view details)

Uploaded Aug 12, 2026 CPython 3.13Windows x86-64

splineax_klujax-0.5.0.post3-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl (3.3 MB view details)

Uploaded Aug 12, 2026 CPython 3.13manylinux: glibc 2.26+ x86-64manylinux: glibc 2.28+ x86-64

splineax_klujax-0.5.0.post3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (3.2 MB view details)

Uploaded Aug 12, 2026 CPython 3.13manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

splineax_klujax-0.5.0.post3-cp313-cp313-macosx_11_0_arm64.whl (277.0 kB view details)

Uploaded Aug 12, 2026 CPython 3.13macOS 11.0+ ARM64

splineax_klujax-0.5.0.post3-cp312-cp312-win_arm64.whl (186.1 kB view details)

Uploaded Aug 12, 2026 CPython 3.12Windows ARM64

splineax_klujax-0.5.0.post3-cp312-cp312-win_amd64.whl (202.9 kB view details)

Uploaded Aug 12, 2026 CPython 3.12Windows x86-64

splineax_klujax-0.5.0.post3-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl (3.3 MB view details)

Uploaded Aug 12, 2026 CPython 3.12manylinux: glibc 2.26+ x86-64manylinux: glibc 2.28+ x86-64

splineax_klujax-0.5.0.post3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (3.2 MB view details)

Uploaded Aug 12, 2026 CPython 3.12manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

splineax_klujax-0.5.0.post3-cp312-cp312-macosx_11_0_arm64.whl (277.0 kB view details)

Uploaded Aug 12, 2026 CPython 3.12macOS 11.0+ ARM64

splineax_klujax-0.5.0.post3-cp311-cp311-win_arm64.whl (187.2 kB view details)

Uploaded Aug 12, 2026 CPython 3.11Windows ARM64

splineax_klujax-0.5.0.post3-cp311-cp311-win_amd64.whl (202.8 kB view details)

Uploaded Aug 12, 2026 CPython 3.11Windows x86-64

splineax_klujax-0.5.0.post3-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl (3.3 MB view details)

Uploaded Aug 12, 2026 CPython 3.11manylinux: glibc 2.26+ x86-64manylinux: glibc 2.28+ x86-64

splineax_klujax-0.5.0.post3-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl (3.2 MB view details)

Uploaded Aug 12, 2026 CPython 3.11manylinux: glibc 2.26+ ARM64manylinux: glibc 2.28+ ARM64

splineax_klujax-0.5.0.post3-cp311-cp311-macosx_11_0_arm64.whl (277.4 kB view details)

Uploaded Aug 12, 2026 CPython 3.11macOS 11.0+ ARM64

File details

Details for the file splineax_klujax-0.5.0.post3-cp314-cp314-win_arm64.whl.

File metadata

Download URL: splineax_klujax-0.5.0.post3-cp314-cp314-win_arm64.whl
Upload date: Aug 12, 2026
Size: 192.4 kB
Tags: CPython 3.14, Windows ARM64
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for splineax_klujax-0.5.0.post3-cp314-cp314-win_arm64.whl
Algorithm	Hash digest
SHA256	`b9e70d0c865b748bde8276a8293753855a403c41cea3565dd083a16a632e4fdd`
MD5	`690e7674cf9040c95ad67ad0cfead7a7`
BLAKE2b-256	`7154a3cbea2b66489dd7c0d0f3680ecaf79a9647e3d512275388254685f29d2d`

Algorithm	Hash digest
SHA256	`4a1ffd4a64691c6b179f2b2f51d5b29adbe3b4458b693accf18ab7753c75be96`
MD5	`201f98fbf2ceae9a3354745b51a70fdb`
BLAKE2b-256	`4a07cd9b9c8ae3008002978686265da8738a04a6bb6674c4175b32ccadb4500b`

Algorithm	Hash digest
SHA256	`e923da07fa0537c0578855e6ee15fb52f2feba13b66ecd84998567977859fb50`
MD5	`aab0237d33798a56ea3af899c99235ad`
BLAKE2b-256	`19bf639232643e05e3f975ea06cc207b2918fc282753a785a99444599e7e7a07`

Algorithm	Hash digest
SHA256	`0ac0fd2265f7d304cf9f4c5f3040faa9480982a6b007e0774b21c252eb97855a`
MD5	`688bee6b4c0978a8e1d8c220a74c8f40`
BLAKE2b-256	`a6fbefe0ba6b6a9a2640b416a47c637a64e01cf4521ced4e67fb51c98d3e42d1`

Algorithm	Hash digest
SHA256	`0fa9f361ab99844f9972c4ce380aa962b37de4dd51c4c01daa399e4f8b4541fc`
MD5	`d8b7ac0021818c4d4f7dee4a244888af`
BLAKE2b-256	`3ca88b8e09380adf868a9bfc0dc1f79f8c2dbe9a6490b580a4a5d7812b9b6cd3`

Algorithm	Hash digest
SHA256	`6a314dcb189aaa863f33f52813177b175a38672d08e7f2a1343633ab41a478a1`
MD5	`4f8803f28fd664a98434399dc79bb4db`
BLAKE2b-256	`a2ad564e781d31e58e64d3bf4d5b9fac5b63c884ed794420e410d982e270fb4f`

Algorithm	Hash digest
SHA256	`d495e8ca13562c6120e14c9c9d038ebd5b507933d7718f583a54dd6d56fcbd94`
MD5	`0c10df9b3c1c10ef2a39253d35724945`
BLAKE2b-256	`91c745f2da869cfa35b19308828b9bee90e84ae4257d9bce02f6eb2fb39c33d5`

Algorithm	Hash digest
SHA256	`05c71ddab757bf9e8b3f692784bbd3c31e0c5f4e9d227f7b495fdea169453952`
MD5	`8e5744a4d69b4358a4a1d77f07740a11`
BLAKE2b-256	`eb66a481a09fdba927d82b445a755c02b7d852293717beb18b276f255964ca92`

Algorithm	Hash digest
SHA256	`30b280bb504a532a7ac1222cb47ff61f227092cda599581653d3d1604adaf28d`
MD5	`8cba9b2208f9c4feedec78a1690f24f1`
BLAKE2b-256	`0f7daa12798fc1f6cca340c5d0a8db05e51f75796fb920f744fefd176ce5cddc`

Algorithm	Hash digest
SHA256	`817be79876a275d36068b311fe8d04b0b2944ec64ff419848f7010fd4b880be0`
MD5	`8ce8f5c9a43b8fa854830917a0a34707`
BLAKE2b-256	`78e66542a30acfb42bf542762c10dd7c7b7835e5d5a55971026860345b449ba5`

Algorithm	Hash digest
SHA256	`1e629273796b9fe07850b962e79b97583749a93d01f55042a0a8edf5995a012d`
MD5	`76e0fc891ecf9c93e140e864de39e325`
BLAKE2b-256	`eb298ea2a384969c2383fd121a0cf3525b34291eeb7829761a478dfcdada27f4`

Algorithm	Hash digest
SHA256	`f37cf62a406cf553e985e3060ae41c240bd8af1cb5c62cb32f3105b609fc4a92`
MD5	`d643d2c3576ac0b93d7708972cb16acf`
BLAKE2b-256	`746a649c66cfe22043bb862297eca43374d3eb11f5008bfafef15754a2d3ba7e`

Algorithm	Hash digest
SHA256	`67b73f0a74db884dd247fa85447f850aaf1f4d45a9c9c86430b7ae115fa120b7`
MD5	`d66ff583e7d43fdab38c3bc1e60a61a0`
BLAKE2b-256	`03ce3ad872286f80a53edcbb30986fbfde6a251e987b67138a3fdeafcae32ede`

Algorithm	Hash digest
SHA256	`46970ed5426b2e5d119b7475557d8b0ace782f29581b3c00526f15038bae425f`
MD5	`86af1af80cd677ae62cfe61524fbe630`
BLAKE2b-256	`3da0a65c0e404d0f41497f90abcf75209f025808cf0ed5c40c93329a9be12d46`

Algorithm	Hash digest
SHA256	`b8a614c473346ae2fbd943b2a9e6281de6b96aff343a30908474dcfba9ba0b3c`
MD5	`7d77d0c1b6828654012839036bb2785f`
BLAKE2b-256	`1764c0f8c001cf1ef8d266dba77a97301e366b5f2f9d11f25c569d821817f421`

Algorithm	Hash digest
SHA256	`aff527fde572b824940cc6ac723f936315a32f8a58319eac4a0a72ac6c919c7e`
MD5	`6cb1d285c8e137f6bcf9c736baa9c8d1`
BLAKE2b-256	`367ce0cbefbc88b3ee59a065d82edfd5d4e87e39585ec101b12c1b7a48255d07`

Algorithm	Hash digest
SHA256	`862d130d4940c172e7a35a9fc168bc125ffe9e859dd9b3ce19361fcdf2770d5a`
MD5	`b855cdc70a8cdb40f84c9e9726e7d3d3`
BLAKE2b-256	`ccf5a66c55110acac7355a3b8a614401816e5d95687486aa36ec0399408f8592`

Algorithm	Hash digest
SHA256	`f9a733c9b56d6d6b67ad52ad561c5d25fbb7970faa0a553ef2774af9e58e7690`
MD5	`6ba163fe89faf9d7838e556358ffc4f4`
BLAKE2b-256	`169ebbc74baf02f5ad0c8e4069b7f61261117faeee4fb5d3a078a52003364ff7`

Algorithm	Hash digest
SHA256	`c6c911ffcc3db1e9d79d59331d6ca414abc2c0ab53dc9b2d65541213c20bfb13`
MD5	`928e77f8f63629db230c45304a0d2d8d`
BLAKE2b-256	`a7819e1211c912dd80b8db5254d86852576c4c8e3704833ecbd2531bcaddd32a`

Algorithm	Hash digest
SHA256	`cf93e0172c9681f7a1f152991a5c7ae2f1020a3d7325ffb225f09400955a854a`
MD5	`7834d9e625683091c1cfa26cb6d1f067`
BLAKE2b-256	`e32bb1b39114fa916f3e3070eba92e8606ab7026a655a038885af5c48700b0d9`

splineax-klujax 0.5.0.post3

KLUJAX

CPU & float64

Basic Usage

Basic Example

Advanced Usage

1. High-Performance Transient Pattern (Reusing Symbolic)

Fine-Grained Control (Numeric Factorization)

Safe Refactorization (Status Codes & Conditioning)

Lifecycle & Pointer Pitfalls

The "Ghost Pointer" Problem inside JIT:

The Fix: Explicit Destruction with Dependencies

Summary of Best Practices

Installation

Installing from Source

Linux

MacOs

Windows

License & Credits

Key dates

1 maintainer

Credits

License expression

Requires

Classifiers

Download files

Source Distributions

Built Distributions

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance

File details

File metadata

File hashes

Provenance