def pivoted_cholesky(matrix, max_rank, diag_rtol=<NUM_LIT>, name=None):
with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>dtype = dtype_util.common_dtype([matrix, diag_rtol],<EOL>preferred_dtype=tf.float32)<EOL>matrix = tf.convert_to_tensor(value=matrix, name='<STR_LIT>', dtype=dtype)<EOL>if tensorshape_util.rank(matrix.shape) is None:<EOL><INDENT>raise NotImplementedError('<STR_LIT>')<EOL><DEDENT>max_rank = tf.convert_to_tensor(<EOL>value=max_rank, name='<STR_LIT>', dtype=tf.int64)<EOL>max_rank = tf.minimum(max_rank,<EOL>prefer_static.shape(matrix, out_type=tf.int64)[-<NUM_LIT:1>])<EOL>diag_rtol = tf.convert_to_tensor(<EOL>value=diag_rtol, dtype=dtype, name='<STR_LIT>')<EOL>matrix_diag = tf.linalg.diag_part(matrix)<EOL>orig_error = tf.reduce_max(input_tensor=matrix_diag, axis=-<NUM_LIT:1>)<EOL>def cond(m, pchol, perm, matrix_diag):<EOL><INDENT>"""<STR_LIT>"""<EOL>del pchol<EOL>del perm<EOL>error = tf.linalg.norm(tensor=matrix_diag, ord=<NUM_LIT:1>, axis=-<NUM_LIT:1>)<EOL>max_err = tf.reduce_max(input_tensor=error / orig_error)<EOL>return (m < max_rank) & (tf.equal(m, <NUM_LIT:0>) | (max_err > diag_rtol))<EOL><DEDENT>batch_dims = tensorshape_util.rank(matrix.shape) - <NUM_LIT:2><EOL>def batch_gather(params, indices, axis=-<NUM_LIT:1>):<EOL><INDENT>return tf.gather(params, indices, axis=axis, batch_dims=batch_dims)<EOL><DEDENT>def body(m, pchol, perm, matrix_diag):<EOL><INDENT>"""<STR_LIT>"""<EOL>permuted_diag = batch_gather(matrix_diag, perm[..., m:])<EOL>maxi = tf.argmax(<EOL>input=permuted_diag, axis=-<NUM_LIT:1>, output_type=tf.int64)[..., tf.newaxis]<EOL>maxval = batch_gather(permuted_diag, maxi)<EOL>maxi = maxi + m<EOL>maxval = maxval[..., <NUM_LIT:0>]<EOL>perm = _swap_m_with_i(perm, m, maxi)<EOL>row = batch_gather(matrix, perm[..., m:m + <NUM_LIT:1>], axis=-<NUM_LIT:2>)<EOL>row = batch_gather(row, perm[..., m + <NUM_LIT:1>:])<EOL>prev_rows = pchol[..., :m, :]<EOL>prev_rows_perm_m_onward = batch_gather(prev_rows, perm[..., m + <NUM_LIT:1>:])<EOL>prev_rows_pivot_col = batch_gather(prev_rows, perm[..., m:m + <NUM_LIT:1>])<EOL>row -= tf.reduce_sum(<EOL>input_tensor=prev_rows_perm_m_onward * prev_rows_pivot_col,<EOL>axis=-<NUM_LIT:2>)[..., tf.newaxis, :]<EOL>pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis]<EOL>row = tf.concat([pivot, row / pivot], axis=-<NUM_LIT:1>)<EOL>paddings = tf.concat([<EOL>tf.zeros([prefer_static.rank(pchol) - <NUM_LIT:1>, <NUM_LIT:2>], dtype=tf.int32),<EOL>[[tf.cast(m, tf.int32), <NUM_LIT:0>]]], axis=<NUM_LIT:0>)<EOL>diag_update = tf.pad(tensor=row**<NUM_LIT:2>, paddings=paddings)[..., <NUM_LIT:0>, :]<EOL>reverse_perm = _invert_permutation(perm)<EOL>matrix_diag -= batch_gather(diag_update, reverse_perm)<EOL>row = tf.pad(tensor=row, paddings=paddings)<EOL>row = batch_gather(row, reverse_perm)<EOL>pchol_shape = pchol.shape<EOL>pchol = tf.concat([pchol[..., :m, :], row, pchol[..., m + <NUM_LIT:1>:, :]],<EOL>axis=-<NUM_LIT:2>)<EOL>tensorshape_util.set_shape(pchol, pchol_shape)<EOL>return m + <NUM_LIT:1>, pchol, perm, matrix_diag<EOL><DEDENT>m = np.int64(<NUM_LIT:0>)<EOL>pchol = tf.zeros_like(matrix[..., :max_rank, :])<EOL>matrix_shape = prefer_static.shape(matrix, out_type=tf.int64)<EOL>perm = tf.broadcast_to(<EOL>prefer_static.range(matrix_shape[-<NUM_LIT:1>]), matrix_shape[:-<NUM_LIT:1>])<EOL>_, pchol, _, _ = tf.while_loop(<EOL>cond=cond, body=body, loop_vars=(m, pchol, perm, matrix_diag))<EOL>pchol = tf.linalg.matrix_transpose(pchol)<EOL>tensorshape_util.set_shape(<EOL>pchol, tensorshape_util.concatenate(matrix_diag.shape, [None]))<EOL>return pchol<EOL><DEDENT>
Computes the (partial) pivoted cholesky decomposition of `matrix`.

    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