Module: tfp.experimental.fastgp.mbcg

Modified Batched Conjugate Gradients (mBCG) for JAX.

Classes

class Array: Array base class for JAX

class SymmetricTridiagonalMatrix: Holds a batch of symmetric, tridiagonal matrices.

Functions

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

safe_stack(...): Like jnp.stack, but handles len == 0 or 1.

tridiagonal_det(...): Return the determinant of a tridiagonal matrix.