nifty.re.num.lanczos module

lanczos_tridiag(mat: Callable[[Array], Array], v: Array, *, order: int, tol: float = 1e-12)

Compute a Lanczos tridiagonal and its orthonormal projection matrix.

The tridiagonal matrix is of shape (order x order) and the stack of vectors has shape (order,) + v.shape. The output is padded with zeros after an early Lanczos breakdown.

stochastic_logdet_from_lanczos(tridiag_stack: ~jax.Array, matrix_shape0: int, func: ~typing.Callable = <PjitFunction of <function log>>, *, tol=1e-14)

Estimate a matrix trace from a stack of Lanczos tridiagonals.

stochastic_lq_logdet(mat: ~jax.Array | ~typing.Callable, order: int, n_samples: int, key: int | ~jax.Array, *, shape0: int | None = None, dtype=None, cmap=<function vmap>)

Estimate a log-determinant with stochastic Lanczos quadrature.