jax.experimental.sparse.linalg.spsolve#

jax.experimental.sparse.linalg.spsolve(data, indices, indptr, b, tol=1e-06, reorder=1)[source]#

A sparse direct solver using QR factorization.

Accepts a sparse matrix in CSR format data, indices, indptr arrays. Currently only the CUDA GPU backend is implemented, the CPU backend will fall back to scipy.sparse.linalg.spsolve. Neither the CPU nor the GPU implementation support batching with vmap.

Parameters:

  • data – An array containing the non-zero entries of the CSR matrix.

  • indices – The column indices of the CSR matrix.

  • indptr – The row pointer array of the CSR matrix.

  • b – The right hand side of the linear system.

  • tol – Tolerance to decide if singular or not. Defaults to 1e-6.

  • reorder – The reordering scheme to use to reduce fill-in. No reordering if reorder=0. Otherwise, symrcm, symamd, or csrmetisnd (reorder=1,2,3), respectively. Defaults to symrcm.

Returns:

An array with the same dtype and size as b representing the solution to the sparse linear system.