Native bridge API

Host-side sparse-direct bridge: SuperLU factorization via SciPy.

A thin, non-differentiable escape hatch to battle-tested sparse LU (SuperLU, through scipy.sparse.linalg.splu()) for general sparse systems that fall outside the structured solvers in solvax.direct. The factorization and triangular solves run on the host CPU, entirely outside the JAX trace machinery — these functions must not be called under jit, vmap, or grad. A guard raises a clear RuntimeError if a traced value is passed; if you need staging, wrap the call in jax.pure_callback() yourself, and for gradients combine with solvax.implicit.linear_solve outside jit.

SciPy is an optional dependency, imported lazily; install it with pip install solvax[native].

References

  • X. S. Li, An Overview of SuperLU, ACM Trans. Math. Softw. 31(3), 302 (2005), DOI 10.1145/1089014.1089017.

class solvax.native.SpluFactorization(matrix)

Reusable SuperLU factorization of a scipy sparse matrix.

Factor once, solve many times:

lu = SpluFactorization(A_csr)
x1 = lu.solve(b1)
x2 = lu.solve(b2)
Variables:

shape – shape of the factored matrix.

solve(b)

Solve A x = b with the stored factors.

Parameters:

b – concrete (non-traced) right-hand side, shape (n,) or (n, n_rhs).

Returns:

The solution as a jax array.

Raises:

RuntimeError – if called with a traced value (under jit/vmap/grad).

Return type:

Array

solvax.native.splu_solve(matrix, b)

One-shot host-side sparse-direct solve of matrix @ x = b.

Convenience wrapper: SpluFactorization then a single solve. For repeated solves with the same matrix, construct the factorization once and reuse it.

Parameters:

  • matrix – scipy sparse matrix (CSR or CSC).

  • b – concrete (non-traced) right-hand side, shape (n,) or (n, n_rhs).

Returns:

The solution as a jax array.

Raises:
Return type:

Array