code: def testXlaCompileBug(self):
      inp = tf.Variable([[2.0, 1.0], [1.0, 2.0]])
      self.evaluate(inp.initializer)
      alt_chol = simple_robustified_cholesky
      alt_chol_nojit = tf.function(
          alt_chol, autograph=False, jit_compile=False
      )
      alt_chol_jit = tf.function(
          alt_chol, autograph=False, jit_compile=True
      )
      answer = np.array([[1.4142135, 0.0], [0.70710677, 1.2247449]])
  
      self.assertAllClose(self.evaluate(alt_chol(inp)), answer)
      self.assertAllClose(self.evaluate(alt_chol_nojit(inp)), answer)
      # TODO(phandu): Enable the test again when the bug is resolved.
      # Bug in tensorflow since 2.15.0-dev20230812,
      # see details at https://github.com/tensorflow/tensorflow/issues/61674
      # self.assertAllClose(self.evaluate(alt_chol_jit(inp)), answer)
      del alt_chol_jit
  
      with tf.GradientTape():
          chol_with_grad = alt_chol(inp)
          chol_nojit_with_grad = alt_chol_nojit(inp)
          # Not supported by TF-XLA (WAI), see b/193584244
          # chol_jit_with_grad = alt_chol_jit(inp)
      self.assertAllClose(self.evaluate(chol_with_grad), answer)
      self.assertAllClose(self.evaluate(chol_nojit_with_grad), answer)
  
      # But wrapping the tape in tf.function should work.
      @tf.function(autograph=False, jit_compile=True)
      def jit_with_grad(mat):
          with tf.GradientTape():
              return alt_chol_jit(mat)
  
      # TODO(phandu): Enable the test again when the bug is resolved.
      # Bug in tensorflow since 2.15.0-dev20230812,
      # see details at https://github.com/tensorflow/tensorflow/issues/61674
      # self.assertAllClose(self.evaluate(jit_with_grad(inp)), answer)
      del jit_with_grad
  
file_path: probability/tensorflow_probability/python/experimental/linalg/no_pivot_ldl_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.linalg.no_pivot_ldl import no_pivot_ldl
  from tensorflow_probability.python.experimental.linalg.no_pivot_ldl import simple_robustified_cholesky
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 8 for calling the NoPivotLDLTest obj's underlying member method code with example usage: obj.testXlaCompileBug() and returns: alt_chol_jit