def compare_KLs(sess, feed_dict, mu, Q_chol, P_chols): mu_gpflow = tf.transpose(mu) if mu.shape.ndims == 2 else mu[:, None] Q_chol_gpflow = Q_chol if Q_chol.shape.ndims == 3 else Q_chol[None, ...] KL_gpflow = sess.run(gauss_kl(q_mu=mu_gpflow, q_sqrt=Q_chol_gpflow, K=None), feed_dict=feed_dict) KL_gpt = sess.run(KL(mu_diff=mu, Q_chol=Q_chol, P_chol=None, P=None), feed_dict=feed_dict) assert_allclose(KL_gpflow, KL_gpt) for P_chol in P_chols: P_ndims = P_chol.shape.ndims P = tf.square(P_chol) if P_ndims == 1 else tf.matmul( P_chol, P_chol, transpose_b=True) KL_gpflow = sess.run(gauss_kl(q_mu=mu_gpflow, q_sqrt=Q_chol_gpflow, K=tf.diag(P) if P_ndims == 1 else P), feed_dict=feed_dict) KL_gpt = sess.run(KL(mu_diff=mu, Q_chol=Q_chol, P_chol=P_chol, P=None), feed_dict=feed_dict) assert_allclose(KL_gpflow, KL_gpt) KL_gpt = sess.run(KL(mu_diff=mu, Q_chol=Q_chol, P_chol=None, P=P), feed_dict=feed_dict) assert_allclose(KL_gpflow, KL_gpt) KL_gpt = sess.run(KL(mu_diff=mu, Q_chol=Q_chol, P_chol=P_chol, P=P), feed_dict=feed_dict) assert_allclose(KL_gpflow, KL_gpt)
def test_kl_k_cholesky(diag): """ Test that passing K or K_cholesky yield the same answer """ q_mu = Datum.mu q_sqrt = Datum.sqrt_diag if diag else Datum.sqrt kl_K = gauss_kl(q_mu, q_sqrt, K=Datum.K) kl_K_chol = gauss_kl(q_mu, q_sqrt, K_cholesky=Datum.K_cholesky) np.testing.assert_allclose(kl_K.numpy(), kl_K_chol.numpy())
def KL(self): """The KL divergence from variational distribution to the prior.""" if self.white: return kullback_leiblers.gauss_kl(self.q_mu, self.q_sqrt[None, :, :], None) else: K = self.kernel(self.inducing_points) K += default_jitter() * tf.eye(self.num_inducing, dtype=K.dtype) return kullback_leiblers.gauss_kl(self.q_mu, self.q_sqrt[None, :, :], K)
def test_kl_k_cholesky(session_tf, mu, sqrt, sqrt_diag, K, K_cholesky, diag): """ Test that passing K or K_cholesky yield the same answer """ kl_K = gauss_kl(mu, sqrt_diag if diag else sqrt, K=K) kl_K_chol = gauss_kl(mu, sqrt_diag if diag else sqrt, K_cholesky=K_cholesky) np.testing.assert_allclose(kl_K.eval(), kl_K_chol.eval())
def test_whitened(session_tf, diag, mu, sqrt_diag, I): """ Check that K=Identity and K=None give same answer """ chol_from_diag = tf.stack([tf.diag(sqrt_diag[:, i]) for i in range(Datum.N)]) # N x M x M s = sqrt_diag if diag else chol_from_diag kl_white = gauss_kl(mu, s) kl_nonwhite = gauss_kl(mu, s, I) np.testing.assert_allclose(kl_white.eval(), kl_nonwhite.eval())
def KL(self): """ The KL divergence from the variational distribution to the prior. q ~ N(\mu, S) :return: KL divergence from q(u) = N(q_mu, q_s) to p(u) ~ N(0, Kuu), independently for each GP """ if self.white: return gauss_kl(self.q_mu, self.q_sqrt, K=None) else: return gauss_kl(self.q_mu, self.q_sqrt, self.MM_Ku_prior)
def test_whitened(session_tf, diag, mu, sqrt_diag, I): """ Check that K=Identity and K=None give same answer """ chol_from_diag = tf.stack( [tf.diag(sqrt_diag[:, i]) for i in range(Datum.N)]) # N x M x M s = sqrt_diag if diag else chol_from_diag kl_white = gauss_kl(mu, s) kl_nonwhite = gauss_kl(mu, s, I) np.testing.assert_allclose(kl_white.eval(), kl_nonwhite.eval())
def test_diags(session_tf, white, mu, sqrt_diag, K): """ The covariance of q(x) can be Cholesky matrices or diagonal matrices. Here we make sure the behaviours overlap. """ # the chols are diagonal matrices, with the same entries as the diag representation. chol_from_diag = tf.stack([tf.diag(sqrt_diag[:, i]) for i in range(Datum.N)]) # N x M x M # run kl_diag = gauss_kl(mu, sqrt_diag, K if white else None) kl_dense = gauss_kl(mu, chol_from_diag, K if white else None) np.testing.assert_allclose(kl_diag.eval(), kl_dense.eval())
def test_unknown_size_inputs(): """ Test for #725 and #734. When the shape of the Gaussian's mean had at least one unknown parameter, `gauss_kl` would blow up. This happened because `tf.size` can only output types `tf.int32` or `tf.int64`. """ mu = np.ones([1, 4], dtype=default_float()) sqrt = np.ones([4, 1, 1], dtype=default_float()) known_shape = gauss_kl(*map(tf.constant, [mu, sqrt])) unknown_shape = gauss_kl(mu, sqrt) np.testing.assert_allclose(known_shape, unknown_shape)
def test_whitened(diag): """ Check that K=Identity and K=None give same answer """ chol_from_diag = tf.stack( [tf.linalg.diag(Datum.sqrt_diag[:, i]) for i in range(Datum.N)] # [N, M, M] ) s = Datum.sqrt_diag if diag else chol_from_diag kl_white = gauss_kl(Datum.mu, s) kl_nonwhite = gauss_kl(Datum.mu, s, Datum.I) np.testing.assert_allclose(kl_white, kl_nonwhite)
def test_diags(session_tf, white, mu, sqrt_diag, K): """ The covariance of q(x) can be Cholesky matrices or diagonal matrices. Here we make sure the behaviours overlap. """ # the chols are diagonal matrices, with the same entries as the diag representation. chol_from_diag = tf.stack( [tf.diag(sqrt_diag[:, i]) for i in range(Datum.N)]) # N x M x M # run kl_diag = gauss_kl(mu, sqrt_diag, K if white else None) kl_dense = gauss_kl(mu, chol_from_diag, K if white else None) np.testing.assert_allclose(kl_diag.eval(), kl_dense.eval())
def test_diags(white): """ The covariance of q(x) can be Cholesky matrices or diagonal matrices. Here we make sure the behaviours overlap. """ # the chols are diagonal matrices, with the same entries as the diag representation. chol_from_diag = tf.stack( [tf.linalg.diag(Datum.sqrt_diag[:, i]) for i in range(Datum.N)] # [N, M, M] ) kl_diag = gauss_kl(Datum.mu, Datum.sqrt_diag, Datum.K if white else None) kl_dense = gauss_kl(Datum.mu, chol_from_diag, Datum.K if white else None) np.testing.assert_allclose(kl_diag, kl_dense)
def KL(self): """ The KL divergence from the variational distribution to the prior :return: KL divergence from N(q_mu, q_sqrt) to N(0, I), independently for each GP """ return gauss_kl(self.q_mu, self.q_sqrt)
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)
def build_prior_KL(self): if self.whiten: K = None else: K = Kuu(self.feature, self.kern, jitter=settings.numerics.jitter_level) # (P x) x M x M return kullback_leiblers.gauss_kl(self.q_mu, self.q_sqrt, K)
def myKL2(self): X = np.array([[1., 2., 3.], [1., 2.1, 3.], [1.1, 2., 3.], [1., 2., 3.1]]) Y = np.array([[1.], [2.], [.2], [3.]]) Z = np.array([[1., 2., 3.], [1.3, 2.2, 3.1]]) A = np.tril(np.random.rand(6, 6)) #"cholesky" of S_M B = np.random.rand(6, 1) #mu_M all_kernels = [ kernels.RBF(3), kernels.RBF(2, lengthscales=3., variance=2.) ] all_Zs, all_mfs = init_linear(X, Z, all_kernels) mylayers = Fully_Coupled_Layers(X, Y, Z, all_kernels, all_mfs, all_Zs, mu_M=B, S_M=A) kl = mylayers.KL() session = get_default_session() kl = session.run(kl) Kmm1 = all_kernels[0].compute_K_symm( all_Zs[0]) + np.eye(Z.shape[0]) * settings.jitter Kmm2 = all_kernels[1].compute_K_symm( all_Zs[1]) + np.eye(all_Zs[1].shape[0]) * settings.jitter K_big = scipy.linalg.block_diag(Kmm1, Kmm1, Kmm2) tfKL = gauss_kl(tf.constant(B), tf.constant(A[np.newaxis]), K=tf.constant(K_big)) sess = tf.Session() return kl, sess.run(tfKL)
def prior_kl(self) -> tf.Tensor: """ The KL divergence from the variational distribution to the prior :return: KL divergence from N(w_mu, w_sqrt) to N(0, I) """ return gauss_kl( self.w_mu[:, None], self.w_sqrt[None] if not self.is_mean_field else self.w_sqrt[:, None], )
def prior_kl(self) -> tf.Tensor: """ Returns the KL divergence ``KL[q(u)∥p(u)]`` from the prior ``p(u) = N(0, I)`` to the variational distribution ``q(u) = N(w_mu, w_sqrt²)``. """ return gauss_kl( self.w_mu[:, None], self.w_sqrt[None] if not self.is_mean_field else self.w_sqrt[:, None], )
def test_unknown_size_inputs(session_tf): """ Test for #725 and #734. When the shape of the Gaussian's mean had at least one unknown parameter, `gauss_kl` would blow up. This happened because `tf.size` can only output types `tf.int32` or `tf.int64`. """ mu_ph = tf.placeholder(settings.float_type, [None, None]) sqrt_ph = tf.placeholder(settings.float_type, [None, None, None]) mu = np.ones([1, 4], dtype=settings.float_type) sqrt = np.ones([4, 1, 1], dtype=settings.float_type) feed_dict = {mu_ph: mu, sqrt_ph: sqrt} known_shape_tf = gauss_kl(*map(tf.constant, [mu, sqrt])) unknown_shape_tf = gauss_kl(mu_ph, sqrt_ph) known_shape = session_tf.run(known_shape_tf) unknown_shape = session_tf.run(unknown_shape_tf, feed_dict=feed_dict) np.testing.assert_allclose(known_shape, unknown_shape)
def test_unknown_size_inputs(session_tf): """ Test for #725 and #734. When the shape of the Gaussian's mean had at least one unknown parameter, `gauss_kl` would blow up. This happened because `tf.size` can only output types `tf.int32` or `tf.int64`. """ mu_ph = tf.placeholder(settings.float_type, [None, None]) sqrt_ph = tf.placeholder(settings.float_type, [None, None, None]) mu = np.ones([1, 4], dtype=settings.float_type) sqrt = np.ones([4, 1, 1], dtype=settings.float_type) feed_dict = {mu_ph: mu, sqrt_ph: sqrt} known_shape_tf = gauss_kl(*map(tf.constant, [mu, sqrt])) unknown_shape_tf = gauss_kl(mu_ph, sqrt_ph) known_shape = session_tf.run(known_shape_tf) unknown_shape = session_tf.run(unknown_shape_tf, feed_dict=feed_dict) np.testing.assert_allclose(known_shape, unknown_shape)
def _build_likelihood(self): X = self.X Y = self.Y num_samples = tf.shape(X)[0] if self.whiten: f_mean, f_var = self._build_predict(X, full_cov=False, full_output_cov=False) KL = gauss_kl(self.q_mu, tf.matrix_band_part(self.q_sqrt, -1, 0)) else: f_mean, f_var, Kzz = self._build_predict(X, full_cov=False, full_output_cov=False, return_Kzz=True) KL = gauss_kl(self.q_mu, tf.matrix_band_part(self.q_sqrt, -1, 0), K=Kzz) # compute variational expectations var_exp = self.likelihood.variational_expectations(f_mean, f_var, Y) # scaling for batch size scale = tf.cast(self.num_data, settings.float_type) / tf.cast(num_samples, settings.float_type) return tf.reduce_sum(var_exp) * scale - KL
def build_prior_KL(self, K): if K: KL = 0. for i, k in enumerate(K): KL += gauss_kl_white(self.q_mu[:,(i*self.offset):((i+1)*self.offset)], self.q_sqrt[(i*self.offset):((i+1)*self.offset),:,:], K=k ) return KL else: return gauss_kl(self.q_mu, self.q_sqrt, K=K)
def test_sumkl_equals_batchkl(session_tf, shared_k, diag, mu, sqrt, sqrt_diag, K_batch, K): """ gauss_kl implicitely performs a sum of KL divergences This test checks that doing the sum outside of the function is equivalent For q(X)=prod q(x_l) and p(X)=prod p(x_l), check that sum KL(q(x_l)||p(x_l)) = KL(q(X)||p(X)) Here, q(X) has covariance L x M x M p(X) has covariance L x M x M ( or M x M ) Here, q(x_i) has covariance 1 x M x M p(x_i) has covariance M x M """ s = sqrt_diag if diag else sqrt kl_batch = gauss_kl(mu,s,K if shared_k else K_batch) kl_sum = [] for n in range(Datum.N): kl_sum.append(gauss_kl(mu[:, n][:,None], # M x 1 sqrt_diag[:, n][:, None] if diag else sqrt[n, :, :][None, :, :], # 1 x M x M or M x 1 K if shared_k else K_batch[n, :, :][None,:,:])) # 1 x M x M or M x M kl_sum =tf.reduce_sum(kl_sum) assert_almost_equal(kl_sum.eval(), kl_batch.eval())
def test_oned(session_tf, white, mu, sqrt, K_batch): """ Check that the KL divergence matches a 1D by-hand calculation. """ m = 0 mu1d = mu[m,:][None,:] # 1 x N s1d = sqrt[:,m,m][:,None,None] # N x 1 x 1 K1d = K_batch[:,m,m][:,None,None] # N x 1 x 1 kl = gauss_kl(mu1d,s1d,K1d if not white else None) kl_tf = tf_kl_1d(tf.reshape(mu1d,(-1,)), # N tf.reshape(s1d,(-1,)), # N None if white else tf.reshape(K1d,(-1,))) # N np.testing.assert_allclose(kl.eval(), kl_tf.eval())
def logp(self, X, Y): """ :param X: latent state (n_samples x T x E) :param Y: observations (n_samples x T x D) :return: variational lower bound on \log P(Y|X) (n_samples x T) """ KL = kullback_leiblers.gauss_kl(self.Umu, self.Ucov_chol, None) # () fmean, fvar = self.conditional( X, add_observation_noise=False ) # (n_samples x T x D) and (n_samples x T x D) var_exp = tf.reduce_sum( self.likelihood.variational_expectations(fmean, fvar, Y), -1) # (n_samples x T) return var_exp - KL / tf.cast(tf.shape(X)[1], gp.settings.float_type)
def build_prior_KL(self): # whitening of priors can be implemented here """ This gives KL divergence between inducing points priors and approximated posteriors KL(q(u_g)||p(u_g)) + KL(q(u_f)||p(u_f)) q(u_f) = N(u_f|u_fm,u_fs) p(u_f) = N(u_f|0,Kfmm) q(u_g) = N(u_g|u_gm,u_gs) p(u_g) = N(u_g|0,Kgmm) """ if self.whiten: if self.q_diag: KL = kullback_leiblers.gauss_kl_white_diag(self.u_fm, self.u_fs_sqrt) + \ kullback_leiblers.gauss_kl_white_diag(self.u_gm, self.u_gs_sqrt) else: KL = kullback_leiblers.gauss_kl_white(self.u_fm, self.u_fs_sqrt) + \ kullback_leiblers.gauss_kl_white(self.u_gm, self.u_gs_sqrt) else: Kfmm = self.kernf.K(self.Zf) + tf.eye( self.num_inducing_f, dtype=float_type) * settings.numerics.jitter_level Kgmm = self.kerng.K(self.Zg) + tf.eye( self.num_inducing_g, dtype=float_type) * settings.numerics.jitter_level if self.q_diag: KL = kullback_leiblers.gauss_kl_diag(self.u_fm, self.u_fs_sqrt, Kfmm) + \ kullback_leiblers.gauss_kl_diag(self.u_gm, self.u_gs_sqrt, Kgmm) else: KL = kullback_leiblers.gauss_kl(self.u_fm, self.u_fs_sqrt, Kfmm) + \ kullback_leiblers.gauss_kl(self.u_gm, self.u_gs_sqrt, Kgmm) return KL
def build_prior_KL(self): if self.whiten: if self.q_diag: KL = kullback_leiblers.gauss_kl_white_diag( self.q_mu, self.q_sqrt) else: KL = kullback_leiblers.gauss_kl_white(self.q_mu, self.q_sqrt) else: K = self.kern.K(self.Z) + tf.eye(self.num_inducing, dtype=float_type) * settings.numerics.jitter_level if self.q_diag: KL = kullback_leiblers.gauss_kl_diag(self.q_mu, self.q_sqrt, K) else: KL = kullback_leiblers.gauss_kl(self.q_mu, self.q_sqrt, K) return KL
def test_oned(session_tf, white, mu, sqrt, K_batch): """ Check that the KL divergence matches a 1D by-hand calculation. """ m = 0 mu1d = mu[m, :][None, :] # 1 x N s1d = sqrt[:, m, m][:, None, None] # N x 1 x 1 K1d = K_batch[:, m, m][:, None, None] # N x 1 x 1 kl = gauss_kl(mu1d, s1d, K1d if not white else None) kl_tf = tf_kl_1d( tf.reshape(mu1d, (-1, )), # N tf.reshape(s1d, (-1, )), # N None if white else tf.reshape(K1d, (-1, ))) # N np.testing.assert_allclose(kl.eval(), kl_tf.eval())
def test_sumkl_equals_batchkl(session_tf, shared_k, diag, mu, sqrt, sqrt_diag, K_batch, K): """ gauss_kl implicitely performs a sum of KL divergences This test checks that doing the sum outside of the function is equivalent For q(X)=prod q(x_l) and p(X)=prod p(x_l), check that sum KL(q(x_l)||p(x_l)) = KL(q(X)||p(X)) Here, q(X) has covariance L x M x M p(X) has covariance L x M x M ( or M x M ) Here, q(x_i) has covariance 1 x M x M p(x_i) has covariance M x M """ s = sqrt_diag if diag else sqrt kl_batch = gauss_kl(mu, s, K if shared_k else K_batch) kl_sum = [] for n in range(Datum.N): kl_sum.append( gauss_kl( mu[:, n][:, None], # M x 1 sqrt_diag[:, n][:, None] if diag else sqrt[n, :, :][None, :, :], # 1 x M x M or M x 1 K if shared_k else K_batch[n, :, :][None, :, :])) # 1 x M x M or M x M kl_sum = tf.reduce_sum(kl_sum) assert_almost_equal(kl_sum.eval(), kl_batch.eval())
def test_sumkl_equals_batchkl(shared_k, diag): """ gauss_kl implicitely performs a sum of KL divergences This test checks that doing the sum outside of the function is equivalent For q(X)=prod q(x_l) and p(X)=prod p(x_l), check that sum KL(q(x_l)||p(x_l)) = KL(q(X)||p(X)) Here, q(X) has covariance [L, M, M] p(X) has covariance [L, M, M] ( or [M, M] ) Here, q(x_i) has covariance [1, M, M] p(x_i) has covariance [M, M] """ s = Datum.sqrt_diag if diag else Datum.sqrt kl_batch = gauss_kl(Datum.mu, s, Datum.K if shared_k else Datum.K_batch) kl_sum = [] for n in range(Datum.N): q_mu_n = Datum.mu[:, n][:, None] # [M, 1] q_sqrt_n = ( Datum.sqrt_diag[:, n][:, None] if diag else Datum.sqrt[n, :, :][None, :, :] ) # [1, M, M] or [M, 1] K_n = Datum.K if shared_k else Datum.K_batch[n, :, :][None, :, :] # [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)
def test_oned(white, dim): """ Check that the KL divergence matches a 1D by-hand calculation. """ mu1d = Datum.mu[dim, :][None, :] # [1, N] s1d = Datum.sqrt[:, dim, dim][:, None, None] # [N, 1, 1] K1d = Datum.K_batch[:, dim, dim][:, None, None] # [N, 1, 1] kl = gauss_kl(mu1d, s1d, K1d if not white else None) kl_1d = compute_kl_1d( tf.reshape(mu1d, (-1,)), # N tf.reshape(s1d, (-1,)), # N None if white else tf.reshape(K1d, (-1,)), ) # N np.testing.assert_allclose(kl, kl_1d)
def test_local_kl_gpflow_consistency(w_dim): num_data = 400 means = np.random.randn(num_data, w_dim) encoder = DirectlyParameterizedNormalDiag(num_data, w_dim, means) lv = LatentVariableLayer(encoder=encoder, prior=_zero_one_normal_prior(w_dim)) posteriors = lv._inference_posteriors( [np.random.randn(num_data, 3), np.random.randn(num_data, 2)]) q_mu = posteriors.parameters["loc"] q_sqrt = posteriors.parameters["scale_diag"] gpflow_local_kls = gauss_kl(q_mu, q_sqrt) tfp_local_kls = tf.reduce_sum(lv._local_kls(posteriors)) np.testing.assert_allclose(tfp_local_kls, gpflow_local_kls, rtol=1e-10)
def build_prior_KL(self, K): return gauss_kl(self.q_mu, self.q_sqrt, K=K)
def prior_kl(self, Kuu): """ KL divergence between p(u) = N(0, Kuu) and q(u) = N(μ, S) """ return kullback_leiblers.gauss_kl(self.q_mu[:, None], self.q_sqrt[None, :, :], Kuu)
def KL(self): Ku = None if self.white else self.Ku return gauss_kl(self.Um, self.Us_sqrt, Ku)