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)