code: def testRetryingCholeskyWithBatchAndXLA(self):
      matrix = tf.convert_to_tensor(
          [
              [[1.0, 0.5], [0.5, 1.0]],
              [[1e-2, 1.002e-2], [1.002e-2, 1e-2]],
              [[3.0, 3.03], [3.03, 3.031]],
          ]
      )
  
      expected_shift = tf.convert_to_tensor(
          [
              [[0.0, 0.0], [0.0, 0.0]],
              [[1e-4, 0.0], [0.0, 1e-4]],
              [[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_shift), matrix
          )
      )
  
      (res, shift), 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, res)
      self.assertAllClose(expected_shift[..., 0, 0], shift)
      self.assertAllEqual(expected_grad, grad)
      self.assertAllEqual(expected_grad_of_grad, grad_of_grad)
  
      # Test value and gradients of XLA-compiled `retrying_cholesky`.
      xla_retrying_cholesky = tf.function(
          lambda x: marginal_fns.retrying_cholesky(x, max_iters=6),
          autograph=False,
          jit_compile=True,
      )
      (res, shift), grad, grad_of_grad = self.evaluate(
          _value_and_grads(xla_retrying_cholesky, matrix, has_aux=True)
      )
      self.assertAllClose(expected, res)
      self.assertAllClose(expected_shift[..., 0, 0], shift)
      self.assertAllClose(expected_grad, grad)
      self.assertAllClose(
          expected_grad_of_grad, grad_of_grad, rtol=1e-3
      )
  
      # Test XLA-compilation of `retrying_cholesky` and its gradients.
      @tf.function(autograph=False, jit_compile=True)
      def xla_retrying_cholesky_with_grads(x):
          return _value_and_grads(
              lambda x: marginal_fns.retrying_cholesky(x, max_iters=6),
              x,
              has_aux=True,
          )
  
      (res, shift), grad, grad_of_grad = self.evaluate(
          xla_retrying_cholesky_with_grads(matrix)
      )
      self.assertAllClose(expected, res)
      self.assertAllClose(expected_shift[..., 0, 0], shift)
      self.assertAllClose(expected_grad, grad)
      self.assertAllClose(
          expected_grad_of_grad, grad_of_grad, rtol=1e-3
      )
  
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 3 for calling the MarginalFnsTest obj's underlying member method code with example usage: obj.testRetryingCholeskyWithBatchAndXLA() and returns: _value_and_grads