jax.experimental.sparse.linalg.spsolve

Contents

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.

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.