def test_batch_diag(self):
root = torch.randn(4, 5, 3)
actual = root.matmul(root.transpose(-1, -2))
actual_diag = torch.cat([
actual[0].diag().unsqueeze(0),
actual[1].diag().unsqueeze(0),
actual[2].diag().unsqueeze(0),
actual[3].diag().unsqueeze(0),
])
res = RootLazyTensor(root)
self.assertTrue(approx_equal(actual_diag, res.diag()))
def _get_covariance(self, x1, x2):
k_ux1 = delazify(self.base_kernel(x1, self.inducing_points))
if torch.equal(x1, x2):
covar = RootLazyTensor(k_ux1.matmul(self._inducing_inv_root))
# Diagonal correction for predictive posterior
correction = (self.base_kernel(x1, x2, diag=True) -
covar.diag()).clamp(0, math.inf)
covar = PsdSumLazyTensor(covar, DiagLazyTensor(correction))
else:
k_ux2 = delazify(self.base_kernel(x2, self.inducing_points))
covar = MatmulLazyTensor(
k_ux1.matmul(self._inducing_inv_root),
k_ux2.matmul(self._inducing_inv_root).transpose(-1, -2))
return covar
def test_diag(self):
root = torch.randn(5, 3)
actual = root.matmul(root.transpose(-1, -2))
res = RootLazyTensor(root)
self.assertTrue(approx_equal(actual.diag(), res.diag()))
