tfp.experimental.fastgp.mbcg.modified_batched_conjugate_gradients

Return A^(-1)B and Lanczos tridiagonal matrices.

tfp.experimental.fastgp.mbcg.modified_batched_conjugate_gradients(
    matrix_matrix_multiplier: Callable[[tfp.experimental.fastgp.fast_gp.Array], tfp.experimental.fastgp.fast_gp.Array],
    B: tfp.experimental.fastgp.fast_gp.Array,
    preconditioner_fn: Callable[[tfp.experimental.fastgp.fast_gp.Array], tfp.experimental.fastgp.fast_gp.Array],
    max_iters: int = 20,
    tolerance: float = 1e-06
) -> Tuple[tfp.experimental.fastgp.fast_gp.Array, tfp.experimental.fastgp.mbcg.SymmetricTridiagonalMatrix]

Based on Algorithm 2 on page 14 of https://arxiv.org/pdf/1809.11165.pdf

matrix_matrix_multiplier A function for left-matrix multiplying by an n x n matrix A, which should be symmetric and positive definite.
B An n x t matrix containing the vectors v_i for which we want A^(-1) v_i.
preconditioner_fn A function that applies an invertible linear transformation to its input, designed to increase the rate of convergence by decreasing the condition number. The preconditioner_fn should Act like left application of an n by n linear operator, i.e. preconditioner_fn(n x m) should have shape n x m. Somewhat amazingly, this algorithm doesn't care if the passed preconditioner is a left, right, or split preconditioner -- the same output will result. Just be sure to pass in full_preconditioner.solve when using a SplitPreconditioner.
max_iters Run conjugate gradients for at most this many iterations. Note that no matter the value of max_iters, the loop will run for at most n iterations.
tolerance Stop early if all the errors have less than this magnitude.

A pair (C, T) where C is the n x t matrix A^(-1) B and T is a SymmetricTridiagonalMatrix where the diag part is of shape (t, max_iters) and the off_diag part is of shape (t, max_iters -1).