Python tf.linalg.tridiagonal_solve用法及代碼示例

求解三對角方程組。

用法

tf.linalg.tridiagonal_solve(
    diagonals, rhs, diagonals_format='compact', transpose_rhs=False,
    conjugate_rhs=False, name=None, partial_pivoting=True, perturb_singular=False
)

輸入可以以各種格式提供：matrix , sequence 和 compact，由 diagonals_format arg 指定。

在 matrix 格式中，diagonals 必須是形狀為 [..., M, M] 的張量，其中兩個 inner-most 維度代表三對角矩陣。三個對角線之外的元素將被忽略。

在sequence 格式中，diagonals 以元組或三個形狀張量的列表的形式提供，[..., N] , [..., M] , [..., N] 分別表示上對角線、對角線和下對角線。 N 可以是 M-1 或 M ；在後一種情況下，上對角線的最後一個元素和下對角線的第一個元素將被忽略。

在 compact 格式中，三個對角線組合成一個形狀為 [..., 3, M] 的張量，最後兩個維度依次包含上對角線、對角線和下對角線。與sequence 格式類似，元素diagonals[..., 0, M-1] 和diagonals[..., 2, 0] 被忽略。

建議使用compact 格式作為性能最佳的格式。如果您需要手動將張量轉換為緊湊格式，請使用 tf.gather_nd 。形狀為 [m, m] 的張量的示例：

rhs = tf.constant([...])
matrix = tf.constant([[...]])
m = matrix.shape[0]
dummy_idx = [0, 0]  # An arbitrary element to use as a dummy
indices = [[[i, i + 1] for i in range(m - 1)] + [dummy_idx],  # Superdiagonal
         [[i, i] for i in range(m)],                          # Diagonal
         [dummy_idx] + [[i + 1, i] for i in range(m - 1)]]    # Subdiagonal
diagonals=tf.gather_nd(matrix, indices)
x = tf.linalg.tridiagonal_solve(diagonals, rhs)

無論 diagonals_format , rhs 是形狀為 [..., M] 還是 [..., M, K] 的張量。後者允許同時求解具有相同 left-hand 邊和 K 個不同 right-hand 邊的 K 個係統。如果 transpose_rhs 設置為 True ，則預期形狀為 [..., M] 或 [..., K, M] 。

批尺寸，表示為 ... ，在 diagonals 和 rhs 中必須相同。

輸出是與 rhs 形狀相同的張量：[..., M] 或 [..., M, K] 。

如果輸入矩陣不可逆，則不保證運算會引發錯誤。 tf.debugging.check_numerics 可應用於輸出以檢測可逆性問題。

注意：對於大批量大小，如果 partial_pivoting=True 或有多個 right-hand 邊(K > 1)，GPU 上的計算可能會很慢。如果出現此問題，請考慮是否可以禁用旋轉並使用 K = 1 ，或者考慮使用 CPU。

在 CPU 上，解決方案是通過高斯消元計算的，有或沒有部分旋轉，具體取決於partial_pivoting範圍。在 GPU 上，使用了 Nvidia 的 cuSPARSE 庫：https://docs.nvidia.com/cuda/cusparse/index.html#gtsv

[1] Nicholas J. Higham (2002)。數值算法的準確性和穩定性：第二版。暹。頁。 175. 國際標準書號 978-0-89871-802-7。