tfp.experimental.substrates.jax.math.linalg.pivoted_cholesky

View source on GitHub

Computes the (partial) pivoted cholesky decomposition of matrix.

tfp.experimental.substrates.jax.math.linalg.pivoted_cholesky(
    matrix,
    max_rank,
    diag_rtol=0.001,
    name=None
)

The pivoted Cholesky is a low rank approximation of the Cholesky decomposition of matrix, i.e. as described in [(Harbrecht et al., 2012)][1]. The currently-worst-approximated diagonal element is selected as the pivot at each iteration. This yields from a [B1...Bn, N, N] shaped matrix a [B1...Bn, N, K] shaped rank-K approximation 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)][2], as matmuls and solves can be cheaply done via the Woodbury matrix identity, as implemented by tf.linalg.LinearOperatorLowRankUpdate.

Args:

  • matrix: Floating point Tensor batch of symmetric, positive definite matrices.
  • max_rank: Scalar int 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.T are each lower than element * diag_rtol, iteration is permitted to terminate early.
  • name: Optional name for the op.

Returns:

  • lr: Low rank pivoted Cholesky approximation of matrix.

References

[1]: H Harbrecht, M Peters, R Schneider. On the low-rank approximation by the pivoted Cholesky decomposition. Applied numerical mathematics, 62(4):428-440, 2012.

[2]: 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