tfp.experimental.fastgp.partial_lanczos.my_tridiagonal_solve

Like jax.linalg.tridiagonal_solve, but works for all sizes.