Reimplement MultitaskKernel.forward_diag to use sub kernel forward_diags
@@ -38,21 +38,24 @@ class MultitaskKernel(Kernel):
in :class:`gpytorch.kernels.Kernel` because we need to take the Kronecker product of the diagonals of the
base data kernel and the task kernel.
"""
- x1 = x1.transpose(-2, -3)
- x2 = x2.transpose(-2, -3)
task_indices = torch.arange(self.n_tasks, device=x1.device).long()
task_indices = task_indices.unsqueeze(-1).unsqueeze(-1)
# These are small because they are vectors, therefore it is safe to evaluate them
- covar_i = self.task_covar_module(task_indices).evaluate().squeeze(-1)
- covar_x = self.data_covar_module.forward(x1, x2)
- if isinstance(covar_x, LazyVariable):
- covar_x = covar_x.evaluate()
- covar_x = covar_x.squeeze(-1)
+ covar_i_diag = self.task_covar_module.forward_diag(task_indices)
+ covar_x_diag = self.data_covar_module.forward_diag(x1, x2)
+
+ if isinstance(covar_x_diag, LazyVariable):
+ covar_x_diag = covar_x_diag.evaluate()
+ if isinstance(covar_i_diag, LazyVariable):
+ covar_i_diag = covar_i_diag.evaluate()
+
+ covar_x_diag = covar_x_diag.squeeze(-1)
+ covar_i_diag = covar_i_diag.squeeze(-1)
# Take the Kronecker product of the two diagonals
- res = KroneckerProductLazyVariable(NonLazyVariable(covar_i), NonLazyVariable(covar_x)).evaluate().unsqueeze(-1)
- return res
+ res = KroneckerProductLazyVariable(NonLazyVariable(covar_i_diag), NonLazyVariable(covar_x_diag)).evaluate()
+ return res.unsqueeze(-1)
def forward(self, x1, x2):
task_indices = torch.arange(self.n_tasks, device=x1.device).long()
