code: def low_rank_cholesky(
      matrix, max_rank, trace_atol=0, trace_rtol=0, name=None
  ):
      """Computes a low-rank approximation to the Cholesky decomposition.
  
      This routine is similar to pivoted_cholesky, but works under JAX, at
      the cost of being slightly less numerically stable.
  
      Args:
        matrix: Floating point `Tensor` batch of symmetric, positive definite
          matrices, or a tf.linalg.LinearOperator.
        max_rank: Scalar `int` `Tensor`, the rank at which to truncate the
          approximation.
        trace_atol: Scalar floating point `Tensor` (same dtype as `matrix`). If
          trace_atol > 0 and trace(matrix - LR * LR^t) < trace_atol, the output
          LR matrix is allowed to be of rank less than max_rank.
        trace_rtol: Scalar floating point `Tensor` (same dtype as `matrix`). If
          trace_rtol > 0 and trace(matrix - LR * LR^t) < trace_rtol * trace(matrix),
          the output LR matrix is allowed to be of rank less than max_rank.
        name: Optional name for the op.
  
      Returns:
        A triplet (LR, r, residual_diag) of
        LR: a matrix such that LR * LR^t is approximately the input matrix.
          If matrix is of shape (b1, ..., bn, m, m), then LR will be of shape
          (b1, ..., bn, m, r) where r <= max_rank.
        r: the rank of LR.  If r is < max_rank, then
          trace(matrix - LR * LR^t) < trace_atol, and
        residual_diag: The diagonal entries of matrix - LR * LR^t.  This is
          returned because together with LR, it is useful for preconditioning
          the input matrix.
      """
      with tf.name_scope(name or "low_rank_cholesky"):
          dtype = dtype_util.common_dtype(
              [matrix, trace_atol, trace_rtol], dtype_hint=tf.float32
          )
          if not isinstance(matrix, tf.linalg.LinearOperator):
              matrix = tf.convert_to_tensor(
                  matrix, name="matrix", dtype=dtype
              )
              matrix = tf.linalg.LinearOperatorFullMatrix(matrix)
  
          mtrace = matrix.trace()
          mrank = tensorshape_util.rank(matrix.shape)
          batch_dims = mrank - 2
  
          def lr_cholesky_cond(i, _, residual_diag):
              """Condition for `tf.while_loop` continuation."""
              residual_trace = tf.math.reduce_sum(
                  residual_diag, axis=-1
              )
              atol_terminate = (trace_atol > 0) & tf.reduce_all(
                  residual_trace < trace_atol
              )
              rtol_terminate = (trace_rtol > 0) & tf.reduce_all(
                  residual_trace < trace_rtol * mtrace
              )
              terminate = atol_terminate | rtol_terminate
              # TODO(thomaswc): Return false even if i == 0 when mtrace == 0.0 to
              # avoid division by zero errors.
              return (i == 0) | ~terminate
  
          def lr_cholesky_body(i, lr, residual_diag):
              # 1. Find the maximum entry of the residual diagonal.
              max_j = tf.argmax(
                  residual_diag, axis=-1, output_type=tf.int64
              )[..., tf.newaxis]
  
              # 2. Construct vector v that kills that diagonal entry and its row & col.
              # v = residual_matrix[max_j, :] / sqrt(residual_matrix[max_j, max_j])
              maxval = tf.gather(
                  residual_diag, max_j, axis=-1, batch_dims=batch_dims
              )[..., 0]
              normalizer = tf.sqrt(maxval)
              if callable(getattr(matrix, "row", None)):
                  matrix_row = tf.squeeze(matrix.row(max_j), axis=-2)
              else:
                  matrix_row = tf.gather(
                      matrix.to_dense(),
                      max_j,
                      axis=-1,
                      batch_dims=batch_dims,
                  )[..., 0]
              # residual_matrix[max_j, :] = matrix_row[max_j, :] - (lr * lr^t)[max_j, :]
              # And (lr * lr^t)[max_j, :] = lr[max_j, :] * lr^t
              lr_row_maxj = tf.gather(
                  lr, max_j, axis=-2, batch_dims=batch_dims
              )
              lr_lrt_row = tf.matmul(lr_row_maxj, lr, transpose_b=True)
              lr_lrt_row = tf.squeeze(lr_lrt_row, axis=-2)
              unnormalized_v = matrix_row - lr_lrt_row
              v = unnormalized_v / normalizer[..., tf.newaxis]
  
              # Mask v so that it is zero in row/columns we've already zerod.
              # We can use the sign of the residual_diag as the mask because the input
              # matrix being positive definite implies that the diag starts off
              # positive, and only becomes zero on the entries that we've chosen
              # in previous iterations.
              v = v * tf.math.sign(residual_diag)
  
              # 3. Add v to lr.
              # Conceptually the same as
              #   new_lr = lr
              #   new_lr[..., i] = v
              # but without using assignment or dynamic slices, both of which don't
              # work under JAX.
              # v[..., tf.newaxis] is of shape (batch1, ..., batchn, m, 1) and
              # the one_hot term is of shape (1, max_rank) so their broadcasted product
              # will be of shape (batch1, ..., batchn, m, max_rank), the same as lr.
              new_lr = (
                  lr
                  + v[..., tf.newaxis]
                  * tf.one_hot(
                      indices=i, depth=max_rank, dtype=matrix.dtype
                  )[tf.newaxis, :]
              )
  
              # 4. Compute the new residual_diag = old_residual_diag - v * v
              new_residual_diag = residual_diag - v * v
  
              # Explicitly set new_residual_diag[max_j] = 0 (both to guarantee we never
              # choose its index again, and to let us use the tf.math.sign of the
              # residual as a mask.)
              n = new_residual_diag.shape[-1]
              oh = tf.one_hot(
                  indices=max_j[..., 0],
                  depth=n,
                  on_value=0.0,
                  off_value=1.0,
                  dtype=new_residual_diag.dtype,
              )
              new_residual_diag = new_residual_diag * oh
  
              return i + 1, new_lr, new_residual_diag
  
          lr = tf.zeros(matrix.shape, dtype=matrix.dtype)[
              ..., :max_rank
          ]
  
          mdiag = matrix.diag_part()
          i, lr, residual_diag = tf.while_loop(
              cond=lr_cholesky_cond,
              body=lr_cholesky_body,
              loop_vars=(0, lr, mdiag),
              maximum_iterations=max_rank,
          )
  
          return lr, i, residual_diag
  
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 6 for calling a global function with example usage: low_rank_cholesky(matrix, max_rank, trace_atol, trace_rtol, name) and returns: lr,  i,  residual_diag, unknown, unknown,  new_lr,  new_residual_diag