View source on GitHub
|
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.