Module: tfp.experimental.fastgp.partial_lanczos

Run Lanczos for m iterations to get a good preconditioner.

Modules

mbcg module: Modified Batched Conjugate Gradients (mBCG) for JAX.

Classes

class Array: Array base class for JAX

Functions

gram_schmidt(...): Return a vector after applying the Gram-Schmidt process to it.

make_lanczos_preconditioner(...): Return a preconditioner as a linear operator.

my_tridiagonal_solve(...): Like jax.linalg.tridiagonal_solve, but works for all sizes.

partial_lanczos(...): Returns orthonormal Q and tridiagonal T such that A ~ Q^t T Q.

psd_solve_multishift(...): Solve (A - shift_k I) x_k = v for PSD A.

reorthogonalize(...)

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