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