We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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).

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gpflow/kullback_leiblers.py`
Content:
```
1 # Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 # -*- coding: utf-8 -*-
16 
17 import tensorflow as tf
18 
19 from .config import default_float, default_jitter
20 from .covariances.kuus import Kuu
21 from .inducing_variables import InducingVariables
22 from .kernels import Kernel
23 from .utilities import Dispatcher, to_default_float
24 
25 prior_kl = Dispatcher("prior_kl")
26 
27 
28 @prior_kl.register(InducingVariables, Kernel, object, object)
29 def _(inducing_variable, kernel, q_mu, q_sqrt, whiten=False):
30     if whiten:
31         return gauss_kl(q_mu, q_sqrt, None)
32     else:
33         K = Kuu(inducing_variable, kernel, jitter=default_jitter())  # [P, M, M] or [M, M]
34         return gauss_kl(q_mu, q_sqrt, K)
35 
36 
37 def gauss_kl(q_mu, q_sqrt, K=None, *, K_cholesky=None):
38     """
39     Compute the KL divergence KL[q || p] between
40 
41           q(x) = N(q_mu, q_sqrt^2)
42     and
43           p(x) = N(0, K)    if K is not None
44           p(x) = N(0, I)    if K is None
45 
46     We assume L multiple independent distributions, given by the columns of
47     q_mu and the first or last dimension of q_sqrt. Returns the *sum* of the
48     divergences.
49 
50     q_mu is a matrix ([M, L]), each column contains a mean.
51 
52     q_sqrt can be a 3D tensor ([L, M, M]), each matrix within is a lower
53         triangular square-root matrix of the covariance of q.
54     q_sqrt can be a matrix ([M, L]), each column represents the diagonal of a
55         square-root matrix of the covariance of q.
56 
57     K is the covariance of p (positive-definite matrix).  The K matrix can be
58     passed either directly as `K`, or as its Cholesky factor, `K_cholesky`.  In
59     either case, it can be a single matrix [M, M], in which case the sum of the
60     L KL divergences is computed by broadcasting, or L different covariances
61     [L, M, M].
62 
63     Note: if no K matrix is given (both `K` and `K_cholesky` are None),
64     `gauss_kl` computes the KL divergence from p(x) = N(0, I) instead.
65     """
66 
67     if (K is not None) and (K_cholesky is not None):
68         raise ValueError(
69             "Ambiguous arguments: gauss_kl() must only be passed one of `K` or `K_cholesky`."
70         )
71 
72     is_white = (K is None) and (K_cholesky is None)
73     is_diag = len(q_sqrt.shape) == 2
74 
75     shape_constraints = [
76         (q_mu, ["M", "L"]),
77         (q_sqrt, (["M", "L"] if is_diag else ["L", "M", "M"])),
78     ]
79     if not is_white:
80         if K is not None:
81             shape_constraints.append((K, (["L", "M", "M"] if len(K.shape) == 3 else ["M", "M"])))
82         else:
83             shape_constraints.append(
84                 (K_cholesky, (["L", "M", "M"] if len(K_cholesky.shape) == 3 else ["M", "M"]))
85             )
86     tf.debugging.assert_shapes(shape_constraints, message="gauss_kl() arguments")
87 
88     M, L = tf.shape(q_mu)[0], tf.shape(q_mu)[1]
89 
90     if is_white:
91         alpha = q_mu  # [M, L]
92     else:
93         if K is not None:
94             Lp = tf.linalg.cholesky(K)  # [L, M, M] or [M, M]
95         elif K_cholesky is not None:
96             Lp = K_cholesky  # [L, M, M] or [M, M]
97 
98         is_batched = len(Lp.shape) == 3
99 
100         q_mu = tf.transpose(q_mu)[:, :, None] if is_batched else q_mu  # [L, M, 1] or [M, L]
101         alpha = tf.linalg.triangular_solve(Lp, q_mu, lower=True)  # [L, M, 1] or [M, L]
102 
103     if is_diag:
104         Lq = Lq_diag = q_sqrt
105         Lq_full = tf.linalg.diag(tf.transpose(q_sqrt))  # [L, M, M]
106     else:
107         Lq = Lq_full = tf.linalg.band_part(q_sqrt, -1, 0)  # force lower triangle # [L, M, M]
108         Lq_diag = tf.linalg.diag_part(Lq)  # [M, L]
109 
110     # Mahalanobis term: μqᵀ Σp⁻¹ μq
111     mahalanobis = tf.reduce_sum(tf.square(alpha))
112 
113     # Constant term: - L * M
114     constant = -to_default_float(tf.size(q_mu, out_type=tf.int64))
115 
116     # Log-determinant of the covariance of q(x):
117     logdet_qcov = tf.reduce_sum(tf.math.log(tf.square(Lq_diag)))
118 
119     # Trace term: tr(Σp⁻¹ Σq)
120     if is_white:
121         trace = tf.reduce_sum(tf.square(Lq))
122     else:
123         if is_diag and not is_batched:
124             # K is [M, M] and q_sqrt is [M, L]: fast specialisation
125             LpT = tf.transpose(Lp)  # [M, M]
126             Lp_inv = tf.linalg.triangular_solve(
127                 Lp, tf.eye(M, dtype=default_float()), lower=True
128             )  # [M, M]
129             K_inv = tf.linalg.diag_part(tf.linalg.triangular_solve(LpT, Lp_inv, lower=False))[
130                 :, None
131             ]  # [M, M] -> [M, 1]
132             trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))
133         else:
134             # TODO: broadcast instead of tile when tf allows -- tf2.1 segfaults
135             # (https://github.com/tensorflow/tensorflow/issues/37584).
136             # See # https://github.com/GPflow/GPflow/issues/1321
137             Lp_full = Lp if is_batched else tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])
138             LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)
139             trace = tf.reduce_sum(tf.square(LpiLq))
140 
141     twoKL = mahalanobis + constant - logdet_qcov + trace
142 
143     # Log-determinant of the covariance of p(x):
144     if not is_white:
145         log_sqdiag_Lp = tf.math.log(tf.square(tf.linalg.diag_part(Lp)))
146         sum_log_sqdiag_Lp = tf.reduce_sum(log_sqdiag_Lp)
147         # If K is [L, M, M], num_latent_gps is no longer implicit, no need to multiply the single kernel logdet
148         scale = 1.0 if is_batched else to_default_float(L)
149         twoKL += scale * sum_log_sqdiag_Lp
150 
151     tf.debugging.assert_shapes([(twoKL, ())], message="gauss_kl() return value")  # returns scalar
152     return 0.5 * twoKL
153
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
             ]  # [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))

{"golden_diff": "diff --git a/gpflow/kullback_leiblers.py b/gpflow/kullback_leiblers.py\n--- a/gpflow/kullback_leiblers.py\n+++ b/gpflow/kullback_leiblers.py\n@@ -15,6 +15,7 @@\n # -*- coding: utf-8 -*-\n \n import tensorflow as tf\n+from packaging.version import Version\n \n from .config import default_float, default_jitter\n from .covariances.kuus import Kuu\n@@ -131,10 +132,11 @@\n             ]  # [M, M] -> [M, 1]\n             trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))\n         else:\n-            # TODO: broadcast instead of tile when tf allows -- tf2.1 segfaults\n-            # (https://github.com/tensorflow/tensorflow/issues/37584).\n-            # See # https://github.com/GPflow/GPflow/issues/1321\n-            Lp_full = Lp if is_batched else tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])\n+            if is_batched or Version(tf.__version__) >= Version(\"2.2\"):\n+                Lp_full = Lp\n+            else:\n+                # workaround for segfaults when broadcasting in TensorFlow<2.2\n+                Lp_full = tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])\n             LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)\n             trace = tf.reduce_sum(tf.square(LpiLq))\n", "issue": "Remove tiling in gauss_kl when tensorflow supports proper broadcasting\nCurrently, 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).\nRemove tiling in gauss_kl when tensorflow supports proper broadcasting\nCurrently, 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).\n", "before_files": [{"content": "# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# -*- coding: utf-8 -*-\n\nimport tensorflow as tf\n\nfrom .config import default_float, default_jitter\nfrom .covariances.kuus import Kuu\nfrom .inducing_variables import InducingVariables\nfrom .kernels import Kernel\nfrom .utilities import Dispatcher, to_default_float\n\nprior_kl = Dispatcher(\"prior_kl\")\n\n\n@prior_kl.register(InducingVariables, Kernel, object, object)\ndef _(inducing_variable, kernel, q_mu, q_sqrt, whiten=False):\n    if whiten:\n        return gauss_kl(q_mu, q_sqrt, None)\n    else:\n        K = Kuu(inducing_variable, kernel, jitter=default_jitter())  # [P, M, M] or [M, M]\n        return gauss_kl(q_mu, q_sqrt, K)\n\n\ndef gauss_kl(q_mu, q_sqrt, K=None, *, K_cholesky=None):\n    \"\"\"\n    Compute the KL divergence KL[q || p] between\n\n          q(x) = N(q_mu, q_sqrt^2)\n    and\n          p(x) = N(0, K)    if K is not None\n          p(x) = N(0, I)    if K is None\n\n    We assume L multiple independent distributions, given by the columns of\n    q_mu and the first or last dimension of q_sqrt. Returns the *sum* of the\n    divergences.\n\n    q_mu is a matrix ([M, L]), each column contains a mean.\n\n    q_sqrt can be a 3D tensor ([L, M, M]), each matrix within is a lower\n        triangular square-root matrix of the covariance of q.\n    q_sqrt can be a matrix ([M, L]), each column represents the diagonal of a\n        square-root matrix of the covariance of q.\n\n    K is the covariance of p (positive-definite matrix).  The K matrix can be\n    passed either directly as `K`, or as its Cholesky factor, `K_cholesky`.  In\n    either case, it can be a single matrix [M, M], in which case the sum of the\n    L KL divergences is computed by broadcasting, or L different covariances\n    [L, M, M].\n\n    Note: if no K matrix is given (both `K` and `K_cholesky` are None),\n    `gauss_kl` computes the KL divergence from p(x) = N(0, I) instead.\n    \"\"\"\n\n    if (K is not None) and (K_cholesky is not None):\n        raise ValueError(\n            \"Ambiguous arguments: gauss_kl() must only be passed one of `K` or `K_cholesky`.\"\n        )\n\n    is_white = (K is None) and (K_cholesky is None)\n    is_diag = len(q_sqrt.shape) == 2\n\n    shape_constraints = [\n        (q_mu, [\"M\", \"L\"]),\n        (q_sqrt, ([\"M\", \"L\"] if is_diag else [\"L\", \"M\", \"M\"])),\n    ]\n    if not is_white:\n        if K is not None:\n            shape_constraints.append((K, ([\"L\", \"M\", \"M\"] if len(K.shape) == 3 else [\"M\", \"M\"])))\n        else:\n            shape_constraints.append(\n                (K_cholesky, ([\"L\", \"M\", \"M\"] if len(K_cholesky.shape) == 3 else [\"M\", \"M\"]))\n            )\n    tf.debugging.assert_shapes(shape_constraints, message=\"gauss_kl() arguments\")\n\n    M, L = tf.shape(q_mu)[0], tf.shape(q_mu)[1]\n\n    if is_white:\n        alpha = q_mu  # [M, L]\n    else:\n        if K is not None:\n            Lp = tf.linalg.cholesky(K)  # [L, M, M] or [M, M]\n        elif K_cholesky is not None:\n            Lp = K_cholesky  # [L, M, M] or [M, M]\n\n        is_batched = len(Lp.shape) == 3\n\n        q_mu = tf.transpose(q_mu)[:, :, None] if is_batched else q_mu  # [L, M, 1] or [M, L]\n        alpha = tf.linalg.triangular_solve(Lp, q_mu, lower=True)  # [L, M, 1] or [M, L]\n\n    if is_diag:\n        Lq = Lq_diag = q_sqrt\n        Lq_full = tf.linalg.diag(tf.transpose(q_sqrt))  # [L, M, M]\n    else:\n        Lq = Lq_full = tf.linalg.band_part(q_sqrt, -1, 0)  # force lower triangle # [L, M, M]\n        Lq_diag = tf.linalg.diag_part(Lq)  # [M, L]\n\n    # Mahalanobis term: \u03bcq\u1d40 \u03a3p\u207b\u00b9 \u03bcq\n    mahalanobis = tf.reduce_sum(tf.square(alpha))\n\n    # Constant term: - L * M\n    constant = -to_default_float(tf.size(q_mu, out_type=tf.int64))\n\n    # Log-determinant of the covariance of q(x):\n    logdet_qcov = tf.reduce_sum(tf.math.log(tf.square(Lq_diag)))\n\n    # Trace term: tr(\u03a3p\u207b\u00b9 \u03a3q)\n    if is_white:\n        trace = tf.reduce_sum(tf.square(Lq))\n    else:\n        if is_diag and not is_batched:\n            # K is [M, M] and q_sqrt is [M, L]: fast specialisation\n            LpT = tf.transpose(Lp)  # [M, M]\n            Lp_inv = tf.linalg.triangular_solve(\n                Lp, tf.eye(M, dtype=default_float()), lower=True\n            )  # [M, M]\n            K_inv = tf.linalg.diag_part(tf.linalg.triangular_solve(LpT, Lp_inv, lower=False))[\n                :, None\n            ]  # [M, M] -> [M, 1]\n            trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))\n        else:\n            # TODO: broadcast instead of tile when tf allows -- tf2.1 segfaults\n            # (https://github.com/tensorflow/tensorflow/issues/37584).\n            # See # https://github.com/GPflow/GPflow/issues/1321\n            Lp_full = Lp if is_batched else tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])\n            LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)\n            trace = tf.reduce_sum(tf.square(LpiLq))\n\n    twoKL = mahalanobis + constant - logdet_qcov + trace\n\n    # Log-determinant of the covariance of p(x):\n    if not is_white:\n        log_sqdiag_Lp = tf.math.log(tf.square(tf.linalg.diag_part(Lp)))\n        sum_log_sqdiag_Lp = tf.reduce_sum(log_sqdiag_Lp)\n        # If K is [L, M, M], num_latent_gps is no longer implicit, no need to multiply the single kernel logdet\n        scale = 1.0 if is_batched else to_default_float(L)\n        twoKL += scale * sum_log_sqdiag_Lp\n\n    tf.debugging.assert_shapes([(twoKL, ())], message=\"gauss_kl() return value\")  # returns scalar\n    return 0.5 * twoKL\n", "path": "gpflow/kullback_leiblers.py"}], "after_files": [{"content": "# Copyright 2016-2020 The GPflow Contributors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# -*- coding: utf-8 -*-\n\nimport tensorflow as tf\nfrom packaging.version import Version\n\nfrom .config import default_float, default_jitter\nfrom .covariances.kuus import Kuu\nfrom .inducing_variables import InducingVariables\nfrom .kernels import Kernel\nfrom .utilities import Dispatcher, to_default_float\n\nprior_kl = Dispatcher(\"prior_kl\")\n\n\n@prior_kl.register(InducingVariables, Kernel, object, object)\ndef _(inducing_variable, kernel, q_mu, q_sqrt, whiten=False):\n    if whiten:\n        return gauss_kl(q_mu, q_sqrt, None)\n    else:\n        K = Kuu(inducing_variable, kernel, jitter=default_jitter())  # [P, M, M] or [M, M]\n        return gauss_kl(q_mu, q_sqrt, K)\n\n\ndef gauss_kl(q_mu, q_sqrt, K=None, *, K_cholesky=None):\n    \"\"\"\n    Compute the KL divergence KL[q || p] between\n\n          q(x) = N(q_mu, q_sqrt^2)\n    and\n          p(x) = N(0, K)    if K is not None\n          p(x) = N(0, I)    if K is None\n\n    We assume L multiple independent distributions, given by the columns of\n    q_mu and the first or last dimension of q_sqrt. Returns the *sum* of the\n    divergences.\n\n    q_mu is a matrix ([M, L]), each column contains a mean.\n\n    q_sqrt can be a 3D tensor ([L, M, M]), each matrix within is a lower\n        triangular square-root matrix of the covariance of q.\n    q_sqrt can be a matrix ([M, L]), each column represents the diagonal of a\n        square-root matrix of the covariance of q.\n\n    K is the covariance of p (positive-definite matrix).  The K matrix can be\n    passed either directly as `K`, or as its Cholesky factor, `K_cholesky`.  In\n    either case, it can be a single matrix [M, M], in which case the sum of the\n    L KL divergences is computed by broadcasting, or L different covariances\n    [L, M, M].\n\n    Note: if no K matrix is given (both `K` and `K_cholesky` are None),\n    `gauss_kl` computes the KL divergence from p(x) = N(0, I) instead.\n    \"\"\"\n\n    if (K is not None) and (K_cholesky is not None):\n        raise ValueError(\n            \"Ambiguous arguments: gauss_kl() must only be passed one of `K` or `K_cholesky`.\"\n        )\n\n    is_white = (K is None) and (K_cholesky is None)\n    is_diag = len(q_sqrt.shape) == 2\n\n    shape_constraints = [\n        (q_mu, [\"M\", \"L\"]),\n        (q_sqrt, ([\"M\", \"L\"] if is_diag else [\"L\", \"M\", \"M\"])),\n    ]\n    if not is_white:\n        if K is not None:\n            shape_constraints.append((K, ([\"L\", \"M\", \"M\"] if len(K.shape) == 3 else [\"M\", \"M\"])))\n        else:\n            shape_constraints.append(\n                (K_cholesky, ([\"L\", \"M\", \"M\"] if len(K_cholesky.shape) == 3 else [\"M\", \"M\"]))\n            )\n    tf.debugging.assert_shapes(shape_constraints, message=\"gauss_kl() arguments\")\n\n    M, L = tf.shape(q_mu)[0], tf.shape(q_mu)[1]\n\n    if is_white:\n        alpha = q_mu  # [M, L]\n    else:\n        if K is not None:\n            Lp = tf.linalg.cholesky(K)  # [L, M, M] or [M, M]\n        elif K_cholesky is not None:\n            Lp = K_cholesky  # [L, M, M] or [M, M]\n\n        is_batched = len(Lp.shape) == 3\n\n        q_mu = tf.transpose(q_mu)[:, :, None] if is_batched else q_mu  # [L, M, 1] or [M, L]\n        alpha = tf.linalg.triangular_solve(Lp, q_mu, lower=True)  # [L, M, 1] or [M, L]\n\n    if is_diag:\n        Lq = Lq_diag = q_sqrt\n        Lq_full = tf.linalg.diag(tf.transpose(q_sqrt))  # [L, M, M]\n    else:\n        Lq = Lq_full = tf.linalg.band_part(q_sqrt, -1, 0)  # force lower triangle # [L, M, M]\n        Lq_diag = tf.linalg.diag_part(Lq)  # [M, L]\n\n    # Mahalanobis term: \u03bcq\u1d40 \u03a3p\u207b\u00b9 \u03bcq\n    mahalanobis = tf.reduce_sum(tf.square(alpha))\n\n    # Constant term: - L * M\n    constant = -to_default_float(tf.size(q_mu, out_type=tf.int64))\n\n    # Log-determinant of the covariance of q(x):\n    logdet_qcov = tf.reduce_sum(tf.math.log(tf.square(Lq_diag)))\n\n    # Trace term: tr(\u03a3p\u207b\u00b9 \u03a3q)\n    if is_white:\n        trace = tf.reduce_sum(tf.square(Lq))\n    else:\n        if is_diag and not is_batched:\n            # K is [M, M] and q_sqrt is [M, L]: fast specialisation\n            LpT = tf.transpose(Lp)  # [M, M]\n            Lp_inv = tf.linalg.triangular_solve(\n                Lp, tf.eye(M, dtype=default_float()), lower=True\n            )  # [M, M]\n            K_inv = tf.linalg.diag_part(tf.linalg.triangular_solve(LpT, Lp_inv, lower=False))[\n                :, None\n            ]  # [M, M] -> [M, 1]\n            trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))\n        else:\n            if is_batched or Version(tf.__version__) >= Version(\"2.2\"):\n                Lp_full = Lp\n            else:\n                # workaround for segfaults when broadcasting in TensorFlow<2.2\n                Lp_full = tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])\n            LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)\n            trace = tf.reduce_sum(tf.square(LpiLq))\n\n    twoKL = mahalanobis + constant - logdet_qcov + trace\n\n    # Log-determinant of the covariance of p(x):\n    if not is_white:\n        log_sqdiag_Lp = tf.math.log(tf.square(tf.linalg.diag_part(Lp)))\n        sum_log_sqdiag_Lp = tf.reduce_sum(log_sqdiag_Lp)\n        # If K is [L, M, M], num_latent_gps is no longer implicit, no need to multiply the single kernel logdet\n        scale = 1.0 if is_batched else to_default_float(L)\n        twoKL += scale * sum_log_sqdiag_Lp\n\n    tf.debugging.assert_shapes([(twoKL, ())], message=\"gauss_kl() return value\")  # returns scalar\n    return 0.5 * twoKL\n", "path": "gpflow/kullback_leiblers.py"}]}