Utils

Utility functions for the linear algebra module.

add_jitter

add_jitter(
    matrix: Array, jitter: float | Array = 1e-06
) -> Array

Add jitter to the diagonal of a matrix for numerical stability.

cholesky_factor

cholesky_factor(
    op: AbstractLinearOperator,
) -> lx.AbstractLinearOperator

Cholesky factor of a PSD operator. Returns lower-triangular L s.t. A = L L^T.

logdet

logdet(op: AbstractLinearOperator) -> jax.Array

Log-determinant of a PSD operator via its Cholesky factor.