Repository: GPflow/GPflow

Issue:
Remove tiling in gauss_kl when tensorflow supports proper broadcasting
Currently, kullback_leiblers.gauss_kl() explicitly tiles as tf.linalg.triangular_solve does not broadcast (worse, it segfaults). Once TensorFlow fixes this issue (https://github.com/tensorflow/tensorflow/issues/37584), we should remove this workaround from our code, as it is very inefficient (lots of unnecessary memory allocation).
Remove tiling in gauss_kl when tensorflow supports proper broadcasting
Currently, kullback_leiblers.gauss_kl() explicitly tiles as tf.linalg.triangular_solve does not broadcast (worse, it segfaults). Once TensorFlow fixes this issue (https://github.com/tensorflow/tensorflow/issues/37584), we should remove this workaround from our code, as it is very inefficient (lots of unnecessary memory allocation).


Discussion:

Resolving patch:
diff --git a/gpflow/kullback_leiblers.py b/gpflow/kullback_leiblers.py
--- a/gpflow/kullback_leiblers.py
+++ b/gpflow/kullback_leiblers.py
@@ -15,6 +15,7 @@
 # -*- coding: utf-8 -*-
 
 import tensorflow as tf
+from packaging.version import Version
 
 from .config import default_float, default_jitter
 from .covariances.kuus import Kuu
@@ -131,10 +132,11 @@ def gauss_kl(q_mu, q_sqrt, K=None, *, K_cholesky=None):
             ]  # [M, M] -> [M, 1]
             trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))
         else:
-            # TODO: broadcast instead of tile when tf allows -- tf2.1 segfaults
-            # (https://github.com/tensorflow/tensorflow/issues/37584).
-            # See # https://github.com/GPflow/GPflow/issues/1321
-            Lp_full = Lp if is_batched else tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])
+            if is_batched or Version(tf.__version__) >= Version("2.2"):
+                Lp_full = Lp
+            else:
+                # workaround for segfaults when broadcasting in TensorFlow<2.2
+                Lp_full = tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])
             LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)
             trace = tf.reduce_sum(tf.square(LpiLq))
 


Test patch:
diff --git a/tests/gpflow/test_kullback_leiblers.py b/tests/gpflow/test_kullback_leiblers.py
--- a/tests/gpflow/test_kullback_leiblers.py
+++ b/tests/gpflow/test_kullback_leiblers.py
@@ -14,6 +14,8 @@
 
 # -*- coding: utf-8 -*-
 
+from unittest.mock import patch
+
 import numpy as np
 import pytest
 import tensorflow as tf
@@ -166,6 +168,25 @@ def test_sumkl_equals_batchkl(shared_k, diag):
     assert_almost_equal(kl_sum, kl_batch)
 
 
+@patch("tensorflow.__version__", "2.1.0")
+def test_sumkl_equals_batchkl_shared_k_not_diag_mocked_tf21():
+    """
+    Version of test_sumkl_equals_batchkl with shared_k=True and diag=False
+    that tests the TensorFlow < 2.2 workaround with tiling still works.
+    """
+    kl_batch = gauss_kl(Datum.mu, Datum.sqrt, Datum.K)
+    kl_sum = []
+    for n in range(Datum.N):
+        q_mu_n = Datum.mu[:, n][:, None]  # [M, 1]
+        q_sqrt_n = Datum.sqrt[n, :, :][None, :, :]  # [1, M, M] or [M, 1]
+        K_n = Datum.K  # [1, M, M] or [M, M]
+        kl_n = gauss_kl(q_mu_n, q_sqrt_n, K=K_n)
+        kl_sum.append(kl_n)
+
+    kl_sum = tf.reduce_sum(kl_sum)
+    assert_almost_equal(kl_sum, kl_batch)
+
+
 @pytest.mark.parametrize("dim", [0, 1])
 @pytest.mark.parametrize("white", [True, False])
 def test_oned(white, dim):