---
id: klujax
version: "0.5.0"
license: LGPL-2.0-only
license_treatment: copyleft
maintenance: active
---

# klujax — a KLU solver for JAX

License: copyleft · Maintenance: active · Downloads: 166.2K/mo

## What it is and what it does

klujax wraps the KLU sparse linear solver from SuiteSparse, enabling efficient solution of sparse linear systems Ax=b within JAX. It accepts sparse matrices in COO format (row indices, column indices, values) and right-hand-side vectors, returning solutions optimized for CPU computation with double precision. The package is designed for scientific computing workflows where sparsity structure is exploited to reduce computation and memory.

The library supports both simple one-shot solving via `solve()` and advanced patterns for high-performance applications. For transient simulations or iterative methods, you can separate the expensive symbolic analysis (sparsity inspection) and numeric factorization (LU decomposition) steps from the solve step, reusing them across multiple systems. This is critical for performance when the matrix structure or values change frequently but the sparsity pattern remains constant.

Use it for:

- Solving sparse Jacobian systems in iterative nonlinear solvers (Newton-Raphson, continuation methods)
- Transient circuit or PDE simulations where the sparsity pattern is fixed but matrix values and RHS change each timestep
- Batched sparse linear solves via jax.vmap over multiple systems with the same structure
- Extracting solutions from sparse finite-element or finite-difference discretizations
- Embedded in JAX-based optimization or inverse problems requiring repeated sparse system solves

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

Solves sparse linear systems in JAX using the KLU algorithm from SuiteSparse, optimized for CPU arrays with float64 or complex128 precision.

Yes, if you need to solve sparse linear systems in JAX and can work within the constraints. The package is actively maintained, has no known vulnerabilities, and offers significant performance advantages over dense solvers for sparse problems. Install friction is moderate due to compiled dependencies, but pre-built wheels cover common platforms. The LGPL-2.0-only license requires attention if you plan proprietary distribution. Best suited for research, scientific computing, and open-source projects.

## Install

pip install klujax
uv add klujax
poetry add klujax

## Installing klujax

Before you install: Medium install friction due to compiled dependencies on SuiteSparse. Pre-built wheels available for Python 3.11+ on Linux, Windows, and macOS (both x86_64 and ARM64). Source builds require SuiteSparse development headers. Package is actively maintained.

License in practice: Licensed under LGPL-2.0-only (copyleft). Derivative works and modifications must be distributed under the same license; proprietary use requires careful licensing review.

Quickstart:

import klujax
import jax.numpy as jnp

A_dense = jnp.array([[2, 3, 0], [3, 0, 4], [0, -1, -3]])
b = jnp.array([8, 45, -3])
Ai, Aj = jnp.where(jnp.abs(A_dense) > 0)
Ax = A_dense[Ai, Aj]
result = klujax.solve(Ai, Aj, Ax, b)

Requires float64 or complex128 arrays (float32/complex64 are cast automatically). Sparse matrix must be coalesced. CPU-only; GPU arrays not supported.

Verify before relying:

- Performance comparison with other sparse solvers (e.g., scipy.sparse.linalg) on typical problem sizes
- Numerical stability characteristics and condition number handling
- Memory overhead for large sparse systems relative to dense solvers

## Package facts

- License: LGPL-2.0-only (copyleft)
- Python support: supports_current
- Install friction: medium
- Maintenance: active
- Downloads: 166.2K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

sparse linear solver JAX, KLU algorithm JAX, sparse matrix solve, efficient sparse linear systems, JAX numerical solver, SuiteSparse KLU wrapper, sparse system solver, sparse-linear-algebra, numerical-computing, jax-integration

[View on SkillFed](https://skillfed.io/packages/klujax) · [View on PyPI](https://pypi.org/project/klujax/)