Use foldl in no_pivot_ldl instead of while_loop.

PiperOrigin-RevId: 385835428

Author: midfield

tensorflow_probability/python/experimental/linalg/BUILD

@@ -83,5 +83,6 @@ py_test(
         "//tensorflow_probability/python/experimental/linalg:no_pivot_ldl",
         "//tensorflow_probability/python/internal:tensorshape_util",
         "//tensorflow_probability/python/internal:test_util",
+#         "//third_party/tensorflow/compiler/jit:xla_cpu_jit",  # DisableOnExport
     ],
 )

tensorflow_probability/python/experimental/linalg/no_pivot_ldl.py

@@ -59,9 +59,11 @@ def no_pivot_ldl(matrix, name='no_pivot_ldl'):
 
   Performs the LDL factorization, using the outer product algorithm from [1]. No
   pivoting (or block pivoting) is done, so this should be less stable than
-  e.g. Bunch-Kaufman sytrf. This is implemented as a tf.while_loop, so should
-  have gradients and be accelerator-friendly, but is not particularly
-  performant.
+  e.g. Bunch-Kaufman sytrf. This is implemented as a tf.foldl, so should have
+  gradients and be accelerator-friendly, but is not particularly performant.
+
+  If compiling with XLA, make sure any surrounding GradientTape is also
+  XLA-compiled (b/193584244).
 
   #### References
   [1]: Gene H. Golub, Charles F. Van Loan. Matrix Computations, 4th ed., 2013.
@@ -83,7 +85,7 @@ def no_pivot_ldl(matrix, name='no_pivot_ldl'):
     # TODO(b/182276317) Deal with dynamic ranks better.
     slix = _Slice2Idx(triangular_factor)
 
-    def body(i, triangular_factor):
+    def fn(triangular_factor, i):
       column_head = triangular_factor[..., i, i, tf.newaxis]
       column_tail = triangular_factor[..., i+1:, i]
       rescaled_tail = column_tail / column_head
@@ -97,12 +99,12 @@ def body(i, triangular_factor):
           tf.linalg.band_part(
               tf.einsum('...i,...j->...ij', column_tail, rescaled_tail),
               num_lower=-1, num_upper=0))
-      return i+1, triangular_factor
+      return triangular_factor
 
-    _, triangular_factor = tf.while_loop(
-        cond=lambda i, _: i < tf.shape(triangular_factor)[-1],
-        body=body,
-        loop_vars=(0, triangular_factor))
+    triangular_factor = tf.foldl(
+        fn=fn,
+        elems=tf.range(tf.shape(triangular_factor)[-1]),
+        initializer=triangular_factor)
 
     diag = tf.linalg.diag_part(triangular_factor)
     triangular_factor = tf.linalg.set_diag(

tensorflow_probability/python/experimental/linalg/no_pivot_ldl_test.py

@@ -80,6 +80,34 @@ def testSimpleIndefinite(self):
     eigv, _ = self.evaluate(tf.linalg.eigh(reconstruct))
     self.assertAllTrue(eigv > 0.)
 
+  def testXlaCompileBug(self):
+    inp = tf.Variable([[2., 1.], [1., 2.]])
+    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.70710677, 1.2247449]])
+
+    self.assertAllClose(self.evaluate(alt_chol(inp)), answer)
+    self.assertAllClose(self.evaluate(alt_chol_nojit(inp)), answer)
+    self.assertAllClose(self.evaluate(alt_chol_jit(inp)), answer)
+
+    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)
+
+    self.assertAllClose(self.evaluate(jit_with_grad(inp)), answer)
+
 
 if __name__ == '__main__':
   tf.test.main()