code: def simple_robustified_cholesky(
      matrix, tol=1e-6, name="simple_robustified_cholesky"
  ):
      """Use `no_pivot_ldl` to robustify a Cholesky factorization.
  
      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.
  
      The algorithm proceeds as follows. The input is factored `A = LDL^T`, and the
      too-small diagonal entries of `D` are increased to the tolerance.  Then `L @
      sqrt(D)` is returned.
  
      This algorithm 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. Generally
      speaking, a modified Cholesky factorization of a symmetric matrix `A` is a
      factorization `P(A+E)P^T = LDL^T`, where `P` is a permutation matrix, `L` is
      unit lower triangular, and `D` is (block) diagonal and positive
      definite. Ideally such an algorithm would ensure the following:
  
      1. If `A` is sufficiently positive definite, `E` is zero.
  
      2. If `F` is the smallest matrix (in Frobenius norm) such that `A + F` is
        positive definite, then `E` is not much larger than `F`.
  
      3. `A + E` is reasonably well-conditioned.
  
      4. It is not too much more expensive than the usual Cholesky factorization.
  
      The references give more sophisticated algorithms to ensure all of the
      above. In the simple case where `A = LDL^T` does not require pivoting,
      `simple_robustified_cholesky` will agree and satisfy the above
      criteria. However, in general it may fail to be stable or satisfy 2 and 3.
  
      #### 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]`.
        tol: Minimum for the diagonal.  Default: 1e-6.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: 'simple_robustified_cholesky'.
  
      Returns:
        triangular_factor: The lower triangular Cholesky factor, modified as above.
          This will have shape `[..., n, n]`. Callers should check for `nans` or
          other evidence of instability.
      """
      with tf.name_scope(name) as name:
          matrix = tf.convert_to_tensor(matrix)
          tol = tf.convert_to_tensor(tol, dtype=matrix.dtype)
          triangular_factor, diag = no_pivot_ldl(matrix)
          new_diag = tf.where(diag < tol, tol, diag)
          return tf.einsum(
              "...ij,...j->...ij",
              triangular_factor,
              tf.math.sqrt(new_diag),
          )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 2 for calling a global function with example usage: simple_robustified_cholesky(matrix, tol, name) and returns: tf