jax.experimental.sparse.linalg.spsolve

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

使用QR分解的稀疏直接求解器。

接受CSR格式稀疏矩阵的 data, indices, indptr 数组。目前仅实现了CUDA GPU后端，CPU后端将回退到 scipy.sparse.linalg.spsolve。CPU和GPU实现均不支持使用 vmap 进行批处理。

参数:

• data – 包含CSR矩阵非零条目的数组。

• indices – CSR 矩阵的列索引。

• indptr – CSR矩阵的行指针数组。

• b – 线性系统的右侧。

• tol – 用于决定是否为单数的容差。默认为 1e-6。

• reorder – 用于减少填充的重排方案。如果 reorder=0，则不进行重排。否则，分别使用 symrcm、symamd 或 csrmetisnd（reorder=1,2,3）。默认为 symrcm。

返回:

一个与 b 具有相同 dtype 和大小的一维数组，表示稀疏线性系统的解。