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 和大小的陣列，表示稀疏線性系統的解。