Passing unit tests now

Author: yyexela

gpytorch/models/exact_gp.py

@@ -317,7 +317,7 @@ def __call__(self, *args, **kwargs):
             if settings.debug().on():
                 if not isinstance(full_output, MultivariateNormal):
                     raise RuntimeError("ExactGP.forward must return a MultivariateNormal")
-            full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
+            full_mean, full_covar = full_output.loc, full_output.lazy_covariance_matrix
 
             # Determine the shape of the joint distribution
             batch_shape = full_output.batch_shape

gpytorch/models/exact_prediction_strategies.py

@@ -151,6 +151,7 @@ def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_
             num_tasks = full_output.event_shape[-1]
             full_mean = full_mean.view(*batch_shape, -1, num_tasks)
             fant_mean = full_mean[..., (num_train // num_tasks) :, :]
+            full_targets = full_targets.view(*target_batch_shape, -1)
         else:
             full_mean = full_mean.view(*batch_shape, -1)
             fant_mean = full_mean[..., num_train:]
@@ -211,8 +212,6 @@ def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_
         new_root = new_lt.root_decomposition().root.to_dense()
         new_covar_cache = new_lt.root_inv_decomposition().root.to_dense()
 
-        full_targets = full_targets.view(*target_batch_shape, -1)
-
         # Expand inputs accordingly if necessary (for fantasies at the same points)
         if full_inputs[0].dim() <= full_targets.dim():
             fant_batch_shape = full_targets.shape[:1]
@@ -276,11 +275,7 @@ def train_shape(self):
 
     def exact_prediction(self, joint_mean, joint_covar):
         # Find the components of the distribution that contain test data
-        if not isinstance(self.train_prior_dist, MultitaskMultivariateNormal):
-            test_mean = joint_mean[..., self.num_train :]
-        else:
-            num_tasks = joint_mean.shape[-1]
-            test_mean = joint_mean[..., (self.num_train // num_tasks) :, :]
+        test_mean = joint_mean[..., self.num_train :]
         # For efficiency - we can make things more efficient
         if joint_covar.size(-1) <= settings.max_eager_kernel_size.value():
             test_covar = joint_covar[..., self.num_train :, :].to_dense()
@@ -307,10 +302,11 @@ def exact_predictive_mean(self, test_mean: Tensor, test_train_covar: LinearOpera
         # NOTE TO FUTURE SELF:
         # You **cannot* use addmv here, because test_train_covar may not actually be a non lazy tensor even for an exact
         # GP, and using addmv requires you to to_dense test_train_covar, which is obviously a huge no-no!
-        if not isinstance(self.train_prior_dist, MultitaskMultivariateNormal):
-            res = (test_train_covar @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
+
+        if len(self.mean_cache.shape) == 4:
+            res = (test_train_covar @ self.mean_cache.squeeze(1).unsqueeze(-1)).squeeze(-1)
         else:
-            res = (test_train_covar.unsqueeze(1) @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
+            res = (test_train_covar @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
         res = res + test_mean
 
         return res