code: def _retrying_cholesky_fwd(matrix, jitter, max_iters, unroll):
      """Forward method for `retrying_cholesky`."""
      if unroll < 1:
          raise ValueError(f"Must unroll >=1 iterations. Got {unroll}.")
      matrix = tf.convert_to_tensor(matrix)
      jitter = tf.convert_to_tensor(jitter, dtype=matrix.dtype)
  
      one = tf.ones([], dtype=matrix.dtype)
      ten = tf.convert_to_tensor(10.0, dtype=matrix.dtype)
  
      def cond(i, _, triangular_factor):
          return (i < max_iters) & tf.reduce_any(
              tf.math.is_nan(triangular_factor[..., 0, 0])
          )
  
      def body(i, shift, triangular_factor):
          for _ in range(unroll):
              triangular_factor = tf.linalg.cholesky(
                  tf.linalg.set_diag(
                      matrix, tf.linalg.diag_part(matrix) + shift
                  )
              )
              shift = shift * tf.where(
                  tf.math.is_nan(triangular_factor[..., :1, 0]),
                  ten,
                  one,
              )
          return [i + unroll, shift, triangular_factor]
  
      triangular_factor = tf.linalg.cholesky(matrix)
      shift = tf.where(
          tf.math.is_nan(triangular_factor[..., :1, 0]), jitter, 0.0
      )
      _, shift, triangular_factor = tf.while_loop(
          cond,
          body,
          loop_vars=[tf.convert_to_tensor(0), shift, triangular_factor],
          maximum_iterations=(max_iters + unroll - 1) // unroll,
      )
      shift = tf.stop_gradient(shift)
      return (triangular_factor, shift[..., 0]), (
          triangular_factor,
          matrix,
          jitter,
          shift,
      )
  
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 4 for calling a global function with example usage: _retrying_cholesky_fwd(matrix, jitter, max_iters, unroll) and returns: triangular_factor,  shift,  triangular_factor,  matrix,  jitter,  shift, unknown, list