code: def retrying_cholesky(
      matrix,
      jitter=None,
      max_iters=5,
      unroll=1,
      name="retrying_cholesky",
  ):
      """Computes a modified Cholesky decomposition for a batch of square matrices.
  
      Given a symmetric matrix `A`, this function attempts to give a factorization
      `A + E = LL^T` where `L` is lower triangular, `LL^T` is positive definite, and
      `E` is small in some suitable sense. This is useful for nearly positive
      definite symmetric matrices that are otherwise numerically difficult to
      Cholesky factor.
  
      In particular, this function first attempts a Cholesky decomposition of
      the input matrix.  If that decomposition fails, exponentially-increasing
      diagonal jitter is added to the matrix until either a Cholesky decomposition
      succeeds or until the maximum specified number of iterations is reached.
  
      This function is similar in spirit to a true modified Cholesky factorization
      ([1], [2]). However, it does not use pivoting or other strategies to ensure
      stability, so may not work well for e.g. ill-conditioned matrices.  Further,
      this function may perform multiple Cholesky factorizations, while a true
      modified Cholesky can be done with only slightly more work than a single
      decomposition.
  
      #### References
  
      [1]: Nicholas Higham. What is a modified Cholesky factorization?
        https://nhigham.com/2020/12/22/what-is-a-modified-cholesky-factorization/
  
      [2]: Sheung Hun Cheng and Nicholas Higham, A Modified Cholesky Algorithm Based
        on a Symmetric Indefinite Factorization, SIAM J. Matrix Anal. Appl. 19(4),
        1097–1110, 1998.
  
      Args:
        matrix: A batch of symmetric square matrices, with shape `[..., n, n]`.
        jitter: Initial jitter to add to the diagonal.  Default: 1e-6, unless
          `matrix.dtype` is float64, in which case the default is 1e-10.
        max_iters: Maximum number of times to retry the Cholesky decomposition
          with larger diagonal jitter.  Default: 5.
        unroll: If specified, unroll >1 iterations per while_loop step. This can be
          advantageous on GPU, to avoid host roundtripping.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: 'retrying_cholesky'.
  
      Returns:
        triangular_factor: A Tensor with shape `[..., n, n]`.  The lower triangular
          Cholesky factor, modified as above.  If the Cholesky decomposition failed
          for a batch member, then all lower triangular entries returned for that
          batch member will be NaN.
        diagonal_shift: A tensor of shape `[...]`. `diag_shift[i]` is the value
          added to the diagonal of `matrix[i]` in computing `triangular_factor[i]`.
      """
      with tf.name_scope(name) as name:
          matrix = tf.convert_to_tensor(matrix)
          if jitter is None:
              jitter = 1e-10 if matrix.dtype == tf.float64 else 1e-6
          jitter = tf.convert_to_tensor(jitter, dtype=matrix.dtype)
          return _retrying_cholesky_custom_gradient(
              matrix, jitter, max_iters, unroll
          )
  
file_path: probability/tensorflow_probability/python/experimental/distributions/marginal_fns.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 8 for calling a global function with example usage: retrying_cholesky(matrix, jitter, max_iters, unroll, name) and returns: _retrying_cholesky_custom_gradient