code: def testRetryingCholeskyFailures(self):
      matrix = tf.convert_to_tensor(
          [
              [[1.0, 0.5], [0.5, 1.0]],
              [[1e4, 1e4 + 1], [1e4 + 1, 1e4]],
              [[3.0, 3.03], [3.03, 3.031]],
          ]
      )
  
      expected_eps = tf.convert_to_tensor(
          [
              [[0.0, 0.0], [0.0, 0.0]],
              [[np.nan, np.nan], [np.nan, np.nan]],
              [[0.1, 0.0], [0.0, 0.1]],
          ]
      )
      expected, expected_grad, expected_grad_of_grad = self.evaluate(
          _value_and_grads(
              lambda x: tf.linalg.cholesky(x + expected_eps), matrix
          )
      )
  
      (res, _), grad, grad_of_grad = self.evaluate(
          _value_and_grads(
              lambda x: marginal_fns.retrying_cholesky(x, max_iters=6),
              matrix,
              has_aux=True,
          )
      )
  
      self.assertAllEqual([expected[0], expected[2]], [res[0], res[2]])
      self.assertAllEqual(
          [expected_grad[0], expected_grad[2]], [grad[0], grad[2]]
      )
      self.assertAllEqual(
          [expected_grad_of_grad[0], expected_grad_of_grad[2]],
          [grad_of_grad[0], grad_of_grad[2]],
      )
  
      # Check that the lower-triangular part of `res[1]` is NaN.
      for i in range(res[1].shape[0]):
          self.assertAllNan(res[1, i, : i + 1])
  
      self.assertAllNan(grad[1])
      self.assertAllNan(grad_of_grad[1])
  
file_path: probability/tensorflow_probability/python/experimental/distributions/marginal_fns_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.distributions import marginal_fns
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 5 for calling the MarginalFnsTest obj's underlying member method code with example usage: obj.testRetryingCholeskyFailures() without return types