klujax

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 on this page — 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

klujax on PyPI

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

Evidence: klujax-0.5.0-cp311-cp311-macosx_11_0_arm64.whl; klujax-0.5.0-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; klujax-0.5.0-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl; klujax-0.5.0-cp311-cp311-win_amd64.whl; klujax-0.5.0-cp311-cp311-win_arm64.whl; klujax-0.5.0-cp312-cp312-macosx_11_0_arm64.whl; klujax-0.5.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; klujax-0.5.0-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl; klujax-0.5.0-cp312-cp312-win_amd64.whl; klujax-0.5.0-cp312-cp312-win_arm64.whl; klujax-0.5.0-cp313-cp313-macosx_11_0_arm64.whl; klujax-0.5.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; klujax-0.5.0-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl; klujax-0.5.0-cp313-cp313-win_amd64.whl; klujax-0.5.0-cp313-cp313-win_arm64.whl; klujax-0.5.0-cp314-cp314-macosx_11_0_arm64.whl; klujax-0.5.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; klujax-0.5.0-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl; klujax-0.5.0-cp314-cp314-win_amd64.whl; klujax-0.5.0-cp314-cp314-win_arm64.whl

Tags

More Mathematics packages