|View source on GitHub|
Computes the (partial) pivoted cholesky decomposition of
tfp.experimental.substrates.jax.math.pivoted_cholesky( matrix, max_rank, diag_rtol=0.001, name=None )
The pivoted Cholesky is a low rank approximation of the Cholesky decomposition
matrix, i.e. as described in [(Harbrecht et al., 2012)]. The
currently-worst-approximated diagonal element is selected as the pivot at each
iteration. This yields from a
[B1...Bn, N, N] shaped
N, K] shaped rank-
lr such that
lr @ lr.T ~= matrix.
Note that, unlike the Cholesky decomposition,
lr is not triangular even in
a rectangular-matrix sense. However, under a permutation it could be made
triangular (it has one more zero in each column as you move to the right).
Such a matrix can be useful as a preconditioner for conjugate gradient
optimization, i.e. as in [(Wang et al. 2019)], as matmuls and solves can be
cheaply done via the Woodbury matrix identity, as implemented by
matrix: Floating point
Tensorbatch of symmetric, positive definite matrices.
Tensor, the rank at which to truncate the approximation.
diag_rtol: Scalar floating point
Tensor(same dtype as
matrix). If the errors of all diagonal elements of
lr @ lr.Tare each lower than
element * diag_rtol, iteration is permitted to terminate early.
name: Optional name for the op.
lr: Low rank pivoted Cholesky approximation of
: H Harbrecht, M Peters, R Schneider. On the low-rank approximation by the pivoted Cholesky decomposition. Applied numerical mathematics, 62(4):428-440, 2012.
: K. A. Wang et al. Exact Gaussian Processes on a Million Data Points. arXiv preprint arXiv:1903.08114, 2019. https://arxiv.org/abs/1903.08114