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 具有相同数据类型和大小的数组，表示稀疏线性系统的解。