jax.lax.linalg.tridiagonal_solve

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

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

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

\[A \, X = B\]

参数:

• dl (Array) – 形状为 [..., m] 的向量批次。A 的下对角线：dl[i] := A[i, i-1]，其中 i 在 [0,m) 中。请注意，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]，其中 i 在 [0,m) 中。请注意，dl[m - 1] = 0。

• b (Array) – 右侧矩阵。

返回:

三对角系统的解 X。

返回类型:

Array