code: def pivoted_cholesky(
      matrix,
      max_rank,
      diag_rtol=1e-3,
      return_pivoting_order=False,
      name=None,
  ):
      """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.
        return_pivoting_order: If `True`, return an `int` `Tensor` indicating the
          pivoting order used to produce `lr` (in addition to `lr`).
        name: Optional name for the op.
  
      Returns:
        lr: Low rank pivoted Cholesky approximation of `matrix`.
        perm: (Optional) pivoting order used to produce `lr`.
  
      #### 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
      """
      with tf.name_scope(name or "pivoted_cholesky"):
          dtype = dtype_util.common_dtype(
              [matrix, diag_rtol], dtype_hint=tf.float32
          )
          if not isinstance(matrix, tf.linalg.LinearOperator):
              matrix = tf.convert_to_tensor(
                  matrix, name="matrix", dtype=dtype
              )
          if tensorshape_util.rank(matrix.shape) is None:
              raise NotImplementedError(
                  "Rank of `matrix` must be known statically"
              )
          if isinstance(matrix, tf.linalg.LinearOperator):
              matrix_shape = tf.cast(matrix.shape_tensor(), tf.int64)
          else:
              matrix_shape = ps.shape(matrix, out_type=tf.int64)
  
          max_rank = tf.convert_to_tensor(
              max_rank, name="max_rank", dtype=tf.int64
          )
          max_rank = tf.minimum(max_rank, matrix_shape[-1])
          diag_rtol = tf.convert_to_tensor(
              diag_rtol, dtype=dtype, name="diag_rtol"
          )
          matrix_diag = tf.linalg.diag_part(matrix)
          # matrix is P.D., therefore all matrix_diag > 0, so we don't need abs.
          orig_error = tf.reduce_max(matrix_diag, axis=-1)
  
          def cond(m, pchol, perm, matrix_diag):
              """Condition for `tf.while_loop` continuation."""
              del pchol
              del perm
              error = tf.linalg.norm(matrix_diag, ord=1, axis=-1)
              max_err = tf.reduce_max(error / orig_error)
              return (m < max_rank) & (
                  tf.equal(m, 0) | (max_err > diag_rtol)
              )
  
          batch_dims = tensorshape_util.rank(matrix.shape) - 2
  
          def batch_gather(params, indices, axis=-1):
              return tf.gather(
                  params, indices, axis=axis, batch_dims=batch_dims
              )
  
          def body(m, pchol, perm, matrix_diag):
              """Body of a single `tf.while_loop` iteration."""
              # Here is roughly a numpy, non-batched version of what's going to happen.
              # (See also Algorithm 1 of Harbrecht et al.)
              # 1: maxi = np.argmax(matrix_diag[perm[m:]]) + m
              # 2: maxval = matrix_diag[perm][maxi]
              # 3: perm[m], perm[maxi] = perm[maxi], perm[m]
              # 4: row = matrix[perm[m]][perm[m + 1:]]
              # 5: row -= np.sum(pchol[:m][perm[m + 1:]] * pchol[:m][perm[m]]], axis=-2)
              # 6: pivot = np.sqrt(maxval); row /= pivot
              # 7: row = np.concatenate([[[pivot]], row], -1)
              # 8: matrix_diag[perm[m:]] -= row**2
              # 9: pchol[m, perm[m:]] = row
  
              # Find the maximal position of the (remaining) permuted diagonal.
              # Steps 1, 2 above.
              permuted_diag = batch_gather(matrix_diag, perm[..., m:])
              maxi = tf.argmax(
                  permuted_diag, axis=-1, output_type=tf.int64
              )[..., tf.newaxis]
              maxval = batch_gather(permuted_diag, maxi)
              maxi = maxi + m
              maxval = maxval[..., 0]
              # Update perm: Swap perm[...,m] with perm[...,maxi]. Step 3 above.
              perm = _swap_m_with_i(perm, m, maxi)
              # Step 4.
              if callable(getattr(matrix, "row", None)):
                  row = matrix.row(perm[..., m])[..., tf.newaxis, :]
              else:
                  row = batch_gather(
                      matrix, perm[..., m : m + 1], axis=-2
                  )
              row = batch_gather(row, perm[..., m + 1 :])
              # Step 5.
              prev_rows = pchol[..., :m, :]
              prev_rows_perm_m_onward = batch_gather(
                  prev_rows, perm[..., m + 1 :]
              )
              prev_rows_pivot_col = batch_gather(
                  prev_rows, perm[..., m : m + 1]
              )
              row -= tf.reduce_sum(
                  prev_rows_perm_m_onward * prev_rows_pivot_col, axis=-2
              )[..., tf.newaxis, :]
              # Step 6.
              pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis]
              # Step 7.
              row = tf.concat([pivot, row / pivot], axis=-1)
              # TODO(b/130899118): Pad grad fails with int64 paddings.
              # Step 8.
              paddings = tf.concat(
                  [
                      tf.zeros([ps.rank(pchol) - 1, 2], dtype=tf.int32),
                      [[tf.cast(m, tf.int32), 0]],
                  ],
                  axis=0,
              )
              diag_update = tf.pad(row**2, paddings=paddings)[
                  ..., 0, :
              ]
              reverse_perm = _invert_permutation(perm)
              matrix_diag = matrix_diag - batch_gather(
                  diag_update, reverse_perm
              )
              # Step 9.
              row = tf.pad(row, paddings=paddings)
              # TODO(bjp): Defer the reverse permutation all-at-once at the end?
              row = batch_gather(row, reverse_perm)
              pchol_shape = pchol.shape
              pchol = tf.concat(
                  [pchol[..., :m, :], row, pchol[..., m + 1 :, :]],
                  axis=-2,
              )
              tensorshape_util.set_shape(pchol, pchol_shape)
              return m + 1, pchol, perm, matrix_diag
  
          m = np.int64(0)
          pchol = tf.zeros(matrix_shape, dtype=matrix.dtype)[
              ..., :max_rank, :
          ]
          perm = tf.broadcast_to(
              ps.range(matrix_shape[-1]), matrix_shape[:-1]
          )
          _, pchol, _, _ = tf.while_loop(
              cond=cond,
              body=body,
              loop_vars=(m, pchol, perm, matrix_diag),
          )
          pchol = tf.linalg.matrix_transpose(pchol)
          tensorshape_util.set_shape(
              pchol,
              tensorshape_util.concatenate(matrix_diag.shape, [None]),
          )
  
          if return_pivoting_order:
              return pchol, perm
          else:
              return pchol
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 5 for calling a global function with example usage: pivoted_cholesky(matrix, max_rank, diag_rtol, return_pivoting_order, name) and returns: unknown, tf, unknown,  pchol,  perm,  matrix_diag, pchol,  perm, pchol