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。