def test_diag(self):
avar = a
bvar = b
cvar = c
kp_lazy_var = KroneckerProductLazyTensor(NonLazyTensor(avar),
NonLazyTensor(bvar),
NonLazyTensor(cvar))
res = kp_lazy_var.diag()
actual = kron(kron(avar, bvar), cvar).diag()
self.assertTrue(approx_equal(res, actual))
avar = a.repeat(3, 1, 1)
bvar = b.repeat(3, 1, 1)
cvar = c.repeat(3, 1, 1)
kp_lazy_var = KroneckerProductLazyTensor(NonLazyTensor(avar),
NonLazyTensor(bvar),
NonLazyTensor(cvar))
res = kp_lazy_var.diag()
actual_mat = kron(kron(avar, bvar), cvar)
actual = torch.stack(
[actual_mat[0].diag(), actual_mat[1].diag(), actual_mat[2].diag()])
self.assertTrue(approx_equal(res, actual))
def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
if last_dim_is_batch:
raise RuntimeError(
"MultitaskKernel does not accept the last_dim_is_batch argument."
)
covar_i = self.task_covar_module.covar_matrix
if len(x1.shape[:-2]):
covar_i = covar_i.repeat(*x1.shape[:-2], 1, 1)
if self.bias_only:
covar_i = lazify(
torch.ones_like(covar_i.evaluate())
) # task covariance now all one so it shares covariance but still
# as multitask mean
covar_x = lazify(self.data_covar_module.forward(x1, x2, **params))
res = KroneckerProductLazyTensor(covar_x, covar_i)
return res.diag() if diag else res
def make_posterior_variances(self, joint_covariance_matrix: LazyTensor) -> Tensor:
r"""
Computes the posterior variances given the data points X. As currently
implemented, it computes another forwards call with the stacked data to get out
the joint covariance across all data points.
"""
# TODO: use the exposed joint covariances from the prediction strategy
data_joint_covariance = joint_covariance_matrix.lazy_tensors[
0
].evaluate_kernel()
num_train = self.train_inputs[0].shape[-2]
test_train_covar = data_joint_covariance[..., num_train:, :num_train]
train_train_covar = data_joint_covariance[..., :num_train, :num_train]
test_test_covar = data_joint_covariance[..., num_train:, num_train:]
full_train_train_covar = KroneckerProductLazyTensor(
train_train_covar, *joint_covariance_matrix.lazy_tensors[1:]
)
full_test_test_covar = KroneckerProductLazyTensor(
test_test_covar, *joint_covariance_matrix.lazy_tensors[1:]
)
full_test_train_covar_list = [test_train_covar] + [
*joint_covariance_matrix.lazy_tensors[1:]
]
train_evals, train_evecs = full_train_train_covar.symeig(eigenvectors=True)
# (\kron \Lambda_i + \sigma^2 I)^{-1}
train_inv_evals = DiagLazyTensor(1.0 / (train_evals + self.likelihood.noise))
# compute K_i S_i \hadamard K_i S_i
test_train_hadamard = KroneckerProductLazyTensor(
*[
lt1.matmul(lt2).evaluate() ** 2
for lt1, lt2 in zip(
full_test_train_covar_list, train_evecs.lazy_tensors
)
]
)
# and compute the column sums of
# (\kron K_i S_i * K_i S_i) \tilde{\Lambda}^{-1}
test_train_pred_covar = test_train_hadamard.matmul(train_inv_evals).sum(dim=-1)
pred_variances = full_test_test_covar.diag() - test_train_pred_covar
return pred_variances