tfp.experimental.fastgp.partial_lanczos.tridiagonal_solve_multishift

Solve (T - shift_k I) x_k = v_norm e_1 for batch of sym. tridiagonal T.

T A batched SymmetricTridiagonalMatrix. T.diag should be of shape (k, n).
shifts A tensor of shape (s) representing s distinct scalar shifts.
v_norm A size (k) batch of vector lengths.

A tensor x of shape (s, k, n) that approximately satisfies (T[k] - shifts[i] I) x[i, j, :] = v_norm[k] e_1