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.