def test_q_sqrt_constraints(inducing_points, kernel, mu, white): """ Test that sending in an unconstrained q_sqrt returns the same conditional evaluation and gradients. This is important to match the behaviour of the KL, which enforces q_sqrt is triangular. """ tril = np.tril(rng.randn(Ln, Nn, Nn)) q_sqrt_constrained = Parameter(tril, transform=triangular()) q_sqrt_unconstrained = Parameter(tril) diff_before_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy() assert_allclose(diff_before_gradient_step, 0) kls = [] for q_sqrt in [q_sqrt_constrained, q_sqrt_unconstrained]: with tf.GradientTape() as tape: kl = prior_kl(inducing_points, kernel, mu, q_sqrt, whiten=white) grad = tape.gradient(kl, q_sqrt.unconstrained_variable) q_sqrt.unconstrained_variable.assign_sub(grad) kls.append(kl) diff_kls_before_gradient_step = kls[0] - kls[1] assert_allclose(diff_kls_before_gradient_step, 0) diff_after_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy() assert_allclose(diff_after_gradient_step, 0)
def KL(self): """The KL divergence from variational distribution to the prior.""" return kullback_leiblers.prior_kl(self.inducing_points, self.kernel, self.q_mu, self.q_sqrt, whiten=self.white)
def prior_kl(self) -> tf.Tensor: r""" Returns the KL divergence ``KL[q(u)∥p(u)]`` from the prior ``p(u)`` to the variational distribution ``q(u)``. If this layer uses the :attr:`whiten`\ ed representation, returns ``KL[q(v)∥p(v)]``. """ return prior_kl( self.inducing_variable, self.kernel, self.q_mu, self.q_sqrt, whiten=self.whiten )
def KL(self): """The KL divergence from variational distribution to the prior.""" KL = 0 for i in range(self.num_outputs): q_mu = tf.expand_dims(self.q_mu[:, i], 1) q_sqrt = tf.expand_dims(self.q_sqrt[i], 0) KL += kullback_leiblers.prior_kl(self.inducing_points[i], self.kernels[i], q_mu, q_sqrt, whiten=self.white) return KL
def prior_kl(self) -> tf.Tensor: return kl.prior_kl(self.inducing_variable, self.kernel, self.q_mu, self.q_sqrt, whiten=self.whiten)