jax.experimental.sparse.linalg.spsolve

jax.experimental.sparse.linalg.spsolve#

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

使用 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 和大小的数组，表示稀疏线性系统的解。