klujax
a KLU solver for JAX
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires float64 or complex128 arrays (float32/complex64 are cast automatically).
- Sparse matrix must be coalesced.
- CPU-only; GPU arrays not supported.
License · maintenance · safety
LGPL-2.0-only (copyleft) — Licensed under LGPL-2.0-only (copyleft). Derivative works and modifications must be distributed under the same license; proprietary use requires careful licensing review.
last release 2026-04-15 (121 days)
0 known vulnerabilities (OSV.dev, 2026-08-14) · 166,195 downloads/mo, #10,505 on PyPI
Alternatives
Verify before relying
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)- 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
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 on it.
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
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.
Requires float64 or complex128 arrays (float32/complex64 are cast automatically). Sparse matrix must be coalesced. CPU-only; GPU arrays not supported.
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)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 the current Python release >=3.11 |
| Install friction | Medium. Platform-specific wheel |
| Runtime dependencies | 4 packagesjaxjaxlibjaxtypingnumpy |
| Maintenance | Actively maintained 121 days since the last release |
| First released | |
| Downloads | 166,195 / month, #10,505 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 3 - AlphaOperating System :: OS IndependentProgramming Language :: Python :: 3Programming Language :: Python :: 3.11Programming Language :: Python :: 3.12Programming Language :: Python :: 3.13Programming Language :: Python :: 3.14Topic :: Scientific/Engineering :: Mathematics |
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
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “sparse linear solver JAX”
- klujaxSolves sparse linear systems in JAX using the KLU algorithm from…
- lineaxLineax solves linear systems and least-squares problems in JAX,…
- pyamgPyAMG provides Algebraic Multigrid (AMG) solvers for efficiently…
Give your agent the search over MCP, or paste the wish link into any chat.
More Mathematics packages
NetworkX provides data structures and algorithms for creating, analyzing, and manipulating graphs and networks, supporting everything from simple undirected graphs to complex directed and weighted networks.
kiwisolver is a Python binding to a fast C++ implementation of the Cassowary constraint solver, enabling you to solve systems of linear constraints and inequalities.
Install it if you need to solve constraint systems; skip it if you only need simple linear algebra.
SymPy is a Python library for symbolic mathematics, performing algebraic manipulation, calculus, equation solving, and mathematical expression simplification without numerical approximation.
ContourPy calculates contours of 2D quadrilateral grids using C++11 algorithms wrapped in Python, offering serial and multithreaded implementations without requiring Matplotlib as a dependency.
PyTorch provides GPU-accelerated tensor computation and automatic differentiation for building and training deep neural networks in Python.
onnxruntime loads and executes Open Neural Network Exchange (ONNX) models with a focus on inference performance across CPUs and accelerators.
Install it if you have ONNX models to run in production or development.
See also lineax · jax · jaxlib · qdldl · jax-cuda12-pjrt · jax-cuda13-pjrt · jax-cuda13-plugin · jax-cuda12-plugin · diffrax · munkres