jax.lax.linalg.tridiagonal_solve

jax.lax.linalg.tridiagonal_solve(dl, d, du, b)

计算三对角线性系统的解。

此函数计算三对角线性系统的解：

\[A . X = B\]

参数:

• dl (Array) – 一批形状为 [..., m] 的向量。A 的下对角线：对于 i 在 [0,m) 范围内，dl[i] := A[i, i-1]。注意 dl[0] = 0。

• d (Array) – 一批形状为 [..., m] 的向量。A 的中间对角线：d[i]  := A[i, i] 对于 i 在 [0,m) 中。

• du (Array) – 一批形状为 [..., m] 的向量。A 的上对角线：du[i] := A[i, i+1] 对于 [0,m) 中的 i。注意 dl[m - 1] = 0。

• b (Array) – 右侧矩阵。

返回:

三对角系统的解 X。

返回类型:

Array