text: # Copyright 2016 James Hensman
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  from __future__ import print_function, absolute_import
  from functools import reduce
  import numpy as np
  import gpflow
  import tensorflow as tf
  from .spectral_covariance import make_Kuu, make_Kuf, make_Kuf_np
  from .kronecker_ops import kvs_dot_vec, kron_vec_apply, kvs_dot_mat, kron_mat_apply, kron
  float_type = gpflow.settings.dtypes.float_type
  
  
  class VGP_1d(gpflow.models.GPModel):
      def __init__(self, X, Y, ms, a, b, kern, likelihood,
                   mean_function=gpflow.mean_functions.Zero()):
          """
          Here we assume the interval is [a,b]
          """
          assert X.shape[1] == 1
          assert isinstance(kern, (gpflow.kernels.Matern12,
                                   gpflow.kernels.Matern32,
                                   gpflow.kernels.Matern52))
          kern = kern
          gpflow.models.GPModel.__init__(self, X, Y, kern,
                                        likelihood, mean_function)
          self.num_data = X.shape[0]
          self.num_latent = Y.shape[1]
          self.a = a
          self.b = b
          self.ms = ms
  
          # initialize variational parameters
          Ncos = self.ms.size
          Nsin = self.ms.size - 1
  
          self.q_mu = gpflow.Param(np.zeros((Ncos + Nsin, 1)))
          pos = gpflow.transforms.positive
          self.q_sqrt = gpflow.Param(np.ones(Ncos + Nsin), pos)
  
      @gpflow.params_as_tensors
      def _build_predict(self, X, full_cov=False):
          # given self.q(v), compute q(f)
  
          Kuf = make_Kuf(self.kern, X, self.a, self.b, self.ms)
          Kuu = make_Kuu(self.kern, self.a, self.b, self.ms)
          KiKuf = Kuu.solve(Kuf)
  
          mu = tf.matmul(tf.transpose(KiKuf), self.q_mu)
          tmp1 = tf.expand_dims(self.q_sqrt, 1) * KiKuf
          if full_cov:
              # Kff
              var = self.kern.K(X)
              var = var + tf.matmul(tf.transpose(tmp1), tmp1)  # Projected variance Kfu Ki S Ki Kuf
              var = var - tf.matmul(tf.transpose(Kuf), KiKuf)  # Qff
              var = tf.expand_dims(var, 2)
  
          else:
              var = self.kern.Kdiag(X)  # Kff
              var = var + tf.reduce_sum(tf.square(tmp1), 0)  # Projected variance Kfu Ki [A + WWT] Ki Kuf
              var = var - tf.reduce_sum(Kuf * KiKuf, 0)  # Qff
              var = tf.reshape(var, (-1, 1))
  
          return mu, var
  
      def _build_KL(self):
          """
          We're working in a 'whitened' representation, so this is the KL between
          q(u) and N(0, 1)
          """
          Kuu = make_Kuu(self.kern, self.a, self.b, self.ms)
          Kim = Kuu.solve(self.q_mu)
          KL = 0.5*tf.squeeze(tf.matmul(tf.transpose(Kim), self.q_mu))  # Mahalanobis term
          KL += 0.5 * Kuu.trace_KiX(tf.diag(tf.square(tf.reshape(self.q_sqrt, [-1]))))
          KL += -0.5*tf.cast(tf.size(self.q_mu), float_type)  # Constant term.
          KL += -0.5*tf.reduce_sum(tf.log(tf.square(self.q_sqrt)))  # Log det Q
          KL += 0.5*Kuu.logdet()  # Log det P
          return KL
  
      @gpflow.params_as_tensors
      def _build_likelihood(self):
          # compute the mean and variance of the latent function
          f_mu, f_var = self.build_predict(self.X, full_cov=False)
  
          E_lik = self.likelihood.variational_expectations(f_mu, f_var, self.Y)
          return tf.reduce_sum(E_lik) - self.build_KL()
  
  
  if __name__ == '__main__':
      from matplotlib import pyplot as plt
  
      np.random.seed(0)
      X = np.random.rand(80, 1)*10 - 5
      X = np.sort(X, axis=0)
      Y = np.cos(3*X) + 2*np.sin(5*X) + np.random.randn(*X.shape)*0.8
      Y = np.exp(Y)
  
      plt.ion()
  
      def plot(m, col='r'):
          plt.figure()
          xtest = np.linspace(-8, 6, 1000)[:, None]
          plt.plot(m.X, m.Y, 'kx')
          mu, var = m.predict_f(xtest)
          plt.plot(xtest, np.exp(mu), col)
          plt.plot(xtest, np.exp(mu + 2*np.sqrt(var)), col+'--')
          plt.plot(xtest, np.exp(mu - 2*np.sqrt(var)), col+'--')
  
      for k in [gpflow.kernels.Matern12, gpflow.kernels.Matern32]:
          m = VGP_1d(X, Y, np.arange(1000), a=-6, b=6,
                     kern=k(1),
                     likelihood=gpflow.likelihoods.Exponential())
          m0 = gpflow.vgp.VGP(X, Y, kern=k(1),
                              likelihood=gpflow.likelihoods.Exponential())
  
          m.optimize()
          m0.optimize()
          plot(m, 'r')
          plot(m0, 'b')
          print(m)
          print(m0)
  
  
  class VGP_additive(gpflow.models.GPModel):
      def __init__(self, X, Y, ms, a, b, kerns, likelihood):
          """
          Here we assume the interval is [a,b]
          """
          assert a.size == b.size == len(kerns) == X.shape[1]
          for kern in kerns:
              assert isinstance(kern, (gpflow.kernels.Matern12,
                                       gpflow.kernels.Matern32,
                                       gpflow.kernels.Matern52))
          mf = gpflow.mean_functions.Zero()
          gpflow.models.GPModel.__init__(self, X, Y, kern=None,
                                        likelihood=likelihood, mean_function=mf)
          self.num_latent = 1  # multiple columns not supported in this version
          self.a = a
          self.b = b
          self.ms = ms
          self.input_dim = X.shape[1]
  
          # initialize variational parameters
          Ms = []
          for kern in kerns:
              Ncos_d = self.ms.size
              Nsin_d = self.ms.size - 1
              if isinstance(kern, gpflow.kernels.Matern12):
                  Ncos_d += 1
              elif isinstance(kern, gpflow.kernels.Matern32):
                  Ncos_d += 1
                  Nsin_d += 1
              else:
                  raise NotImplementedError
              Ms.append(Ncos_d + Nsin_d)
  
          self.kerns = gpflow.ParamList(kerns)
  
          self.q_mu = gpflow.Param(np.zeros((np.sum(Ms), 1)))
          pos = gpflow.transforms.positive
          self.q_sqrt = gpflow.ParamList([gpflow.Param(np.ones(M), pos) for M in Ms])
  
      @gpflow.params_as_tensors
      def _build_predict(self, X, full_cov=False):
          # given self.q(v), compute q(f)
  
          Kuf = [make_Kuf(X[:, i:i+1], a, b, self.ms) for i, (a, b) in enumerate(zip(self.a, self.b))]
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          KiKuf = [Kuu_d.solve(Kuf_d) for Kuu_d, Kuf_d in zip(Kuu, Kuf)]
  
          RKiKuf = [Kuu_d.matmul_sqrt(KiKuf_d) for Kuu_d, KiKuf_d in zip(Kuu, KiKuf)]
          KfuKiR = [tf.transpose(RKiKuf_d) for RKiKuf_d in RKiKuf]
  
          mu_d = [tf.matmul(KfuKiR_d, q_mu_d) for KfuKiR_d, q_mu_d in zip(KfuKiR, tf.split(0, self.input_dim, self.q_mu))]
  
          mu = reduce(lambda a, b: a+b, mu_d)
  
          tmp1 = [tf.expand_dims(q_sqrt_d, 1) * RKiKuf_d for q_sqrt_d, RKiKuf_d in zip(self.q_sqrt, RKiKuf)]
          if full_cov:
              raise NotImplementedError
          else:
              # Kff:
              var = reduce(tf.add, [kern.Kdiag(X[:, i:i+1]) for i, kern in enumerate(self.kerns)])
  
              # Projected variance Kfu Ki [A + WWT] Ki Kuf
              var = var + reduce(tf.add, [tf.reduce_sum(tf.square(tmp1_d), 0) for tmp1_d in tmp1])
  
              # Qff
              var = var - reduce(tf.add, [tf.reduce_sum(Kuf_d * KiKuf_d, 0) for Kuf_d, KiKuf_d in zip(Kuf, KiKuf)])
  
              var = tf.reshape(var, (-1, 1))
  
          return mu, var
  
      def build_KL(self):
          """
          We're working in a 'whitened' representation, so this is the KL between
          q(u) and N(0, 1)
          """
          KL = 0.5*tf.reduce_sum(tf.square(self.q_mu))  # Mahalanobis term
          KL += -0.5*tf.cast(tf.size(self.q_mu), float_type)  # Constant term.
          KL += -0.5*reduce(tf.add, [tf.reduce_sum(tf.log(tf.square(q_sqrt_d))) for q_sqrt_d in self.q_sqrt])  # Log det
          KL += 0.5*reduce(tf.add, [tf.reduce_sum(tf.square(q_sqrt_d)) for q_sqrt_d in self.q_sqrt])  # Trace term
  
          return KL
  
      @gpflow.params_as_tensors
      def _build_likelihood(self):
          # compute the mean and variance of the latent function
          f_mu, f_var = self.build_predict(self.X, full_cov=False)
  
          E_lik = self.likelihood.variational_expectations(f_mu, f_var, self.Y)
          return tf.reduce_sum(E_lik) - self.build_KL()
  
  
  class VGP_kron(gpflow.models.GPModel):
      def __init__(self, X, Y, ms, a, b, kerns, likelihood, use_two_krons=False, use_extra_ranks=0):
          """
          Here we assume the interval is [a,b]
          """
          assert a.size == b.size == len(kerns) == X.shape[1]
          for kern in kerns:
              assert isinstance(kern, (gpflow.kernels.Matern12,
                                       gpflow.kernels.Matern32,
                                       gpflow.kernels.Matern52))
          mf = gpflow.mean_functions.Zero()
          gpflow.models.GPModel.__init__(self, X, Y, kern=None,
                                        likelihood=likelihood, mean_function=mf)
          self.num_latent = 1  # multiple columns not supported in this version
          self.a = a
          self.b = b
          self.ms = ms
  
          self.kerns = gpflow.ParamList(kerns)
  
          # initialize variational parameters
          self.Ms = []
          for kern in kerns:
              Ncos_d = self.ms.size
              Nsin_d = self.ms.size - 1
              self.Ms.append(Ncos_d + Nsin_d)
  
          self.q_mu = gpflow.Param(np.zeros((np.prod(self.Ms), 1)))
  
          # The covariance matrix
          self.q_sqrt_kron = gpflow.ParamList([gpflow.Param(np.eye(M)) for M in self.Ms])
          self.use_two_krons = use_two_krons
          self.use_extra_ranks = use_extra_ranks
          assert not (use_extra_ranks and use_two_krons), "can only use one extra covariance structure at a time!"
          if use_two_krons:
              # same as above, but with different init to break symmetry
              self.q_sqrt_kron_2 = gpflow.ParamList([gpflow.Param(np.eye(M)+0.01) for M in self.Ms])
          elif use_extra_ranks:
              self.q_sqrt_W = gpflow.Param(np.zeros((np.prod(self.Ms), use_extra_ranks)))
  
          # pre-compute Kuf
          self._Kuf = [make_Kuf_np(X[:, i:i+1], ai, bi, self.ms)
                 for i, (ai, bi) in enumerate(zip(self.a, self.b))]
  
  
      def __getstate__(self):
          d = gpflow.models.Model.__getstate__(self)
          d.pop('_Kuf')
          return d
  
      def __setstate__(self, d):
          gpflow.models.Model.__setstate__(self, d)
          self._Kuf = [tf.constant(make_Kuf_np(self.X.value[:, i:i+1], ai, bi, self.ms))
                       for i, (ai, bi) in enumerate(zip(self.a, self.b))]
  
      @gpflow.params_as_tensors
      def _build_predict(self, X, full_cov=False):
          # given self.q(v), compute q(f)
  
          Kuf = [make_Kuf(k, X[:, i:i+1], a, b, self.ms) for i, (k, a, b) in enumerate(zip(self.kerns, self.a, self.b))]
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          KiKuf = [Kuu_d.solve(Kuf_d) for Kuu_d, Kuf_d in zip(Kuu, Kuf)]
          KfuKi = [tf.transpose(mat) for mat in KiKuf]
  
          mu = kvs_dot_vec(KfuKi, self.q_mu)
  
          if full_cov:
              raise NotImplementedError
          else:
              # Kff:
              var = reduce(tf.multiply, [k.Kdiag(X[:, i:i+1]) for i, k in enumerate(self.kerns)])
  
              # Projected variance Kfu Ki [WWT] Ki Kuf
              Ls = [tf.matrix_band_part(q_sqrt_d, -1, 0) for q_sqrt_d in self.q_sqrt_kron]
              tmp = [tf.matmul(tf.transpose(L), KiKuf_d) for L, KiKuf_d in zip(Ls, KiKuf)]
              var = var + reduce(tf.multiply, [tf.reduce_sum(tf.square(tmp_d), 0) for tmp_d in tmp])
  
              if self.use_two_krons:
                  Ls = [tf.matrix_band_part(q_sqrt_d, -1, 0) for q_sqrt_d in self.q_sqrt_kron_2]
                  tmp = [tf.matmul(tf.transpose(L), KiKuf_d) for L, KiKuf_d in zip(Ls, KiKuf)]
                  var = var + reduce(tf.multiply, [tf.reduce_sum(tf.square(tmp_d), 0) for tmp_d in tmp])
              elif self.use_extra_ranks:
                  for i in range(self.use_extra_ranks):
                      tmp = kvs_dot_vec(KfuKi, self.q_sqrt_W[:, i:i+1])
                      var = var + tf.reduce_sum(tf.square(tmp), 1)
  
              # Qff
              var = var - reduce(tf.multiply, [tf.reduce_sum(Kuf_d * KiKuf_d, 0) for Kuf_d, KiKuf_d in zip(Kuf, KiKuf)])
  
              var = tf.reshape(var, (-1, 1))
  
          return mu, var
  
      def _build_predict_train(self):
          Kuf = self._Kuf
  
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          KiKuf = [Kuu_d.solve(Kuf_d) for Kuu_d, Kuf_d in zip(Kuu, Kuf)]
          KfuKi = [tf.transpose(mat) for mat in KiKuf]
  
          mu = kvs_dot_vec(KfuKi, self.q_mu)
  
          # Kff:
          var = reduce(tf.multiply, [k.Kdiag(self.X[:, i:i+1]) for i, k in enumerate(self.kerns)])
  
          # Projected variance Kfu Ki [WWT] Ki Kuf
          Ls = [tf.matrix_band_part(q_sqrt_d, -1, 0) for q_sqrt_d in self.q_sqrt_kron]
          tmp = [tf.matmul(tf.transpose(L), KiKuf_d) for L, KiKuf_d in zip(Ls, KiKuf)]
          var = var + reduce(tf.multiply, [tf.reduce_sum(tf.square(tmp_d), 0) for tmp_d in tmp])
  
          if self.use_two_krons:
              Ls = [tf.matrix_band_part(q_sqrt_d, -1, 0) for q_sqrt_d in self.q_sqrt_kron_2]
              tmp = [tf.matmul(tf.transpose(L), KiKuf_d) for L, KiKuf_d in zip(Ls, KiKuf)]
              var = var + reduce(tf.multiply, [tf.reduce_sum(tf.square(tmp_d), 0) for tmp_d in tmp])
          elif self.use_extra_ranks:
              for i in range(self.use_extra_ranks):
                  tmp = kvs_dot_vec(KfuKi, self.q_sqrt_W[:, i:i+1])
                  var = var + tf.reduce_sum(tf.square(tmp), 1)
  
          # Qff
          var = var - reduce(tf.multiply, [tf.reduce_sum(Kuf_d * KiKuf_d, 0) for Kuf_d, KiKuf_d in zip(Kuf, KiKuf)])
  
          return mu, tf.reshape(var, [-1, 1])
  
      @gpflow.autoflow()
      def compute_KL(self):
          return self.build_KL()
  
      def build_KL(self):
          """
          The covariance of q(u) has a kronecker structure, so
          appropriate reductions apply for the trace and logdet terms.
          """
          # Mahalanobis term, m^T K^{-1} m
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          Kim = kron_vec_apply(Kuu, self.q_mu, 'solve')
          KL = 0.5*tf.reduce_sum(self.q_mu * Kim)
  
          # Constant term
          KL += -0.5*tf.cast(tf.size(self.q_mu), float_type)
  
          # Log det term
          Ls = [tf.matrix_band_part(q_sqrt_d, -1, 0) for q_sqrt_d in self.q_sqrt_kron]
          N_others = [float(np.prod(self.Ms)) / M for M in self.Ms]
          Q_logdets = [tf.reduce_sum(tf.log(tf.square(tf.diag_part(L)))) for L in Ls]
          KL += -0.5 * reduce(tf.add, [N*logdet for N, logdet in zip(N_others, Q_logdets)])
  
          # trace term tr(K^{-1} Sigma_q)
          Ss = [tf.matmul(L, tf.transpose(L)) for L in Ls]
          traces = [K.trace_KiX(S) for K, S, in zip(Kuu, Ss)]
          KL += 0.5 * reduce(tf.multiply, traces)  # kron-trace is the produce of traces
  
          # log det term Kuu
          Kuu_logdets = [K.logdet() for K in Kuu]
          KL += 0.5 * reduce(tf.add, [N*logdet for N, logdet in zip(N_others, Kuu_logdets)])
  
          if self.use_two_krons:
              # extra logdet terms:
              Ls_2 = [tf.matrix_band_part(q_sqrt_d, -1, 0) for q_sqrt_d in self.q_sqrt_kron_2]
              LiL = [tf.matrix_triangular_solve(L1, L2) for L1, L2 in zip(Ls, Ls_2)]
              eigvals = [tf.self_adjoint_eig(tf.matmul(tf.transpose(mat), mat))[0] for mat in LiL]  # discard eigenvectors
              eigvals_kronned = kron([tf.reshape(e, [1, -1]) for e in eigvals])
              KL += -0.5 * tf.reduce_sum(tf.log(1 + eigvals_kronned))
  
              # extra trace terms
              Ss = [tf.matmul(L, tf.transpose(L)) for L in Ls_2]
              traces = [K.trace_KiX(S) for K, S, in zip(Kuu, Ss)]
              KL += 0.5 * reduce(tf.multiply, traces)  # kron-trace is the produce of traces
  
          elif self.use_extra_ranks:
              # extra logdet terms
              KiW = kron_mat_apply(Kuu, self.q_sqrt_W, 'solve', self.use_extra_ranks)
              WTKiW = tf.matmul(tf.transpose(self.q_sqrt_W), KiW)
              L_extra = tf.cholesky(np.eye(self.use_extra_ranks) + WTKiW)
              KL += -0.5 * tf.reduce_sum(tf.log(tf.square(tf.diag_part(L_extra))))
  
              # extra trace terms
              KL += 0.5 * tf.reduce_sum(tf.diag_part(WTKiW))
  
          return KL
  
      @gpflow.params_as_tensors
      def _build_likelihood(self):
          # compute the mean and variance of the latent function
          f_mu, f_var = self._build_predict_train()
  
          E_lik = self.likelihood.variational_expectations(f_mu, f_var, self.Y)
          return tf.reduce_sum(E_lik) - self.build_KL()
  
  
  class VGP_kron_anyvar(gpflow.models.GPModel):
      def __init__(self, X, Y, ms, a, b, kerns, likelihood):
          """
          Here we assume the interval is [a,b]
          We do *not* assume that the variance of q(u) has a kronecker structure.
  
          This can get very computationally heavy very quickly, use with caution!.
          """
          assert a.size == b.size == len(kerns) == X.shape[1]
          for kern in kerns:
              assert isinstance(kern, (gpflow.kernels.Matern12,
                                       gpflow.kernels.Matern32,
                                       gpflow.kernels.Matern52))
          mf = gpflow.mean_functions.Zero()
          gpflow.models.GPModel.__init__(self, X, Y, kern=None,
                                        likelihood=likelihood, mean_function=mf)
          self.num_latent = 1  # multiple columns not supported in this version
          self.a = a
          self.b = b
          self.ms = ms
  
          # initialize variational parameters
          Ms = []
          for kern in kerns:
              Ncos_d = self.ms.size
              Nsin_d = self.ms.size - 1
              Ms.append(Ncos_d + Nsin_d)
          self.Ms = Ms
  
          self.kerns = gpflow.ParamList(kerns)
  
          self.q_mu = gpflow.Param(np.zeros((np.prod(Ms), 1)))
  
          # The covariance matrix gets very big very quickly
          self.q_sqrt = gpflow.Param(np.eye(np.prod(Ms)))
  
          # pre-compute the Kuf matrices
          self._Kuf = [tf.constant(make_Kuf_np(X[:, i:i+1], ai, bi, self.ms))
                       for i, (ai, bi) in enumerate(zip(self.a, self.b))]
  
      def __getstate__(self):
          d = gpflow.models.Model.__getstate__(self)
          d.pop('_Kuf')
          return d
  
      def __setstate__(self, d):
          gpflow.models.Model.__setstate__(self, d)
          self._Kuf = [tf.constant(make_Kuf_np(self.X.value[:, i:i+1], ai, bi, self.ms))
                       for i, (ai, bi) in enumerate(zip(self.a, self.b))]
  
      @gpflow.params_as_tensors
      def _build_predict(self, X, full_cov=False):
          # given self.q(v), compute q(f)
  
          Kuf = [make_Kuf(k, X[:, i:i+1], a, b, self.ms) for i, (k, a, b) in enumerate(zip(self.kerns, self.a, self.b))]
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          KiKuf = [Kuu_d.solve(Kuf_d) for Kuu_d, Kuf_d in zip(Kuu, Kuf)]
          KfuKi = [tf.transpose(mat) for mat in KiKuf]
  
          mu = kvs_dot_vec(KfuKi, self.q_mu)
  
          L = tf.matrix_band_part(self.q_sqrt, -1, 0)
          tmp1 = kvs_dot_mat(KfuKi, L, np.prod(self.Ms))
  
          if full_cov:
              raise NotImplementedError
          else:
              # Kff:
              var = reduce(tf.multiply, [k.Kdiag(X[:, i:i+1]) for i, k in enumerate(self.kerns)])
  
              # Projected variance Kfu Ki [WWT] Ki Kuf
              # var = var + reduce(tf.multiply, [tf.reduce_sum(tf.square(tmp1_d), 0) for tmp1_d in tmp1])
              var = var + tf.reduce_sum(tf.square(tmp1), 1)
  
              # Qff
              var = var - reduce(tf.multiply, [tf.reduce_sum(Kuf_d * KiKuf_d, 0) for Kuf_d, KiKuf_d in zip(Kuf, KiKuf)])
  
              var = tf.reshape(var, (-1, 1))
  
          return mu, var
  
      def _build_predict_train(self):
          Kuf = self._Kuf
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          KiKuf = [Kuu_d.solve(Kuf_d) for Kuu_d, Kuf_d in zip(Kuu, Kuf)]
          KfuKi = [tf.transpose(mat) for mat in KiKuf]
  
          mu = kvs_dot_vec(KfuKi, self.q_mu)
          L = tf.matrix_band_part(self.q_sqrt, -1, 0)
          tmp1 = kvs_dot_mat(KfuKi, L, num_cols=np.prod(self.Ms))
  
          # Kff:
          var = reduce(tf.multiply, [k.Kdiag(self.X[:, i:i+1]) for i, k in enumerate(self.kerns)])
          # Projected variance Kfu Ki [WWT] Ki Kuf
          # var = var + reduce(tf.multiply, [tf.reduce_sum(tf.square(tmp1_d), 0) for tmp1_d in tmp1])
          var = var + tf.reduce_sum(tf.square(tmp1), 1)
          # Qff
          var = var - reduce(tf.multiply, [tf.reduce_sum(Kuf_d * KiKuf_d, 0) for Kuf_d, KiKuf_d in zip(Kuf, KiKuf)])
          var = tf.reshape(var, (-1, 1))
  
          return mu, var
  
      @gpflow.autoflow()
      def compute_KL(self):
          return self.build_KL()
  
      def build_KL(self):
          """
          The covariance of q(u) has a kronecker structure, so
          appropriate reductions apply for the trace and logdet terms.
          """
          Kuu = [make_Kuu(kern, a, b, self.ms) for kern, a, b, in zip(self.kerns, self.a, self.b)]
          Kim = kron_vec_apply(Kuu, self.q_mu, 'solve')
          KL = 0.5*tf.reduce_sum(self.q_mu * Kim)  # Mahalanobis term
          KL += -0.5*tf.cast(tf.size(self.q_mu), float_type)  # Constant term.
          L = tf.matrix_band_part(self.q_sqrt, -1, 0)
          Q_logdet = tf.reduce_sum(tf.log(tf.square(tf.diag_part(L))))
          KL += -0.5 * Q_logdet  # log determinat Sigma_q
          S = tf.matmul(L, tf.transpose(L))
          KL += 0.5 * tf.reduce_sum(tf.diag_part(kron_mat_apply(Kuu, S, 'solve', np.prod(self.Ms))))
          Kuu_logdets = [K.logdet() for K in Kuu]
          N_others = [tf.cast(tf.size(self.q_mu) / M, float_type) for M in self.Ms]
          KL += 0.5 * reduce(tf.add, [N*logdet for N, logdet in zip(N_others, Kuu_logdets)])  # kron-logdet P
          return KL
  
      @gpflow.params_as_tensors
      def _build_likelihood(self):
          # compute the mean and variance of the latent function
          f_mu, f_var = self._build_predict_train()
  
          E_lik = self.likelihood.variational_expectations(f_mu, f_var, self.Y)
          return tf.reduce_sum(E_lik) - self.build_KL()
  