sparsax

Lead

sparsax exposes SuiteSparse’s CHOLMOD (Cholesky, for symmetric positive definite matrices) and KLU (LU, for general matrices) to JAX as XLA custom calls. A sparse factorization therefore runs at native speed inside jax.jit, lax.scan, and lax.fori_loop, with no Python callback and no round trip between device and host on each iteration.