Fix unit tests failues caused by upstream breaking changes (#1343)

Summary:
cornellius-gp/gpytorch#2082 changed the shape of the mean constant (and did some other things). These changes fix the resulting unit test breakages.

Pull Request resolved: #1343

Reviewed By: saitcakmak

Differential Revision: D38527688

Pulled By: Balandat

fbshipit-source-id: a9878b3ff5d045b18c3330f4e0707a3ee4be2553

Author: Balandat

botorch/models/gp_regression.py

@@ -156,7 +156,7 @@ def __init__(
             )
             self._subset_batch_dict = {
                 "likelihood.noise_covar.raw_noise": -2,
-                "mean_module.constant": -2,
+                "mean_module.raw_constant": -1,
                 "covar_module.raw_outputscale": -1,
                 "covar_module.base_kernel.raw_lengthscale": -3,
             }
@@ -271,7 +271,7 @@ def __init__(
                 outputscale_prior=GammaPrior(2.0, 0.15),
             )
             self._subset_batch_dict = {
-                "mean_module.constant": -2,
+                "mean_module.raw_constant": -1,
                 "covar_module.raw_outputscale": -1,
                 "covar_module.base_kernel.raw_lengthscale": -3,
             }

botorch/models/gp_regression_fidelity.py

@@ -131,7 +131,7 @@ def __init__(
         )
         self._subset_batch_dict = {
             "likelihood.noise_covar.raw_noise": -2,
-            "mean_module.constant": -2,
+            "mean_module.raw_constant": -1,
             "covar_module.raw_outputscale": -1,
             **subset_batch_dict,
         }
@@ -240,7 +240,7 @@ def __init__(
         )
         self._subset_batch_dict = {
             "likelihood.noise_covar.raw_noise": -2,
-            "mean_module.constant": -2,
+            "mean_module.raw_constant": -1,
             "covar_module.raw_outputscale": -1,
             **subset_batch_dict,
         }

test/acquisition/test_monte_carlo.py

@@ -443,8 +443,8 @@ def test_cache_root(self):
             "likelihood.noise_covar.raw_noise": torch.tensor(
                 [[0.0895], [0.2594]], dtype=torch.float64
             ),
-            "mean_module.constant": torch.tensor(
-                [[-0.4545], [-0.1285]], dtype=torch.float64
+            "mean_module.raw_constant": torch.tensor(
+                [-0.4545, -0.1285], dtype=torch.float64
             ),
             "covar_module.raw_outputscale": torch.tensor(
                 [1.4876, 1.4897], dtype=torch.float64

test/models/test_contextual.py

@@ -17,7 +17,7 @@
 from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
 
 
-class ContextualGPTest(BotorchTestCase):
+class TestContextualGP(BotorchTestCase):
     def test_SACGP(self):
         for dtype in (torch.float, torch.double):
             train_X = torch.tensor(
@@ -45,7 +45,7 @@ def test_SACGP(self):
             num_of_lengthscales = 0
             num_of_outputscales = 0
             for param_name, param in model.named_parameters():
-                if param_name == "mean_module.constant":
+                if param_name == "mean_module.raw_constant":
                     num_of_mean += param.data.shape.numel()
                 elif "raw_lengthscale" in param_name:
                     num_of_lengthscales += param.data.shape.numel()

test/models/test_fully_bayesian.py

@@ -96,7 +96,7 @@ def _get_mcmc_samples(
         mcmc_samples = {
             "lengthscale": torch.rand(num_samples, 1, dim, **tkwargs),
             "outputscale": torch.rand(num_samples, **tkwargs),
-            "mean": torch.randn(num_samples, 1, **tkwargs),
+            "mean": torch.randn(num_samples, **tkwargs),
         }
         if infer_noise:
             mcmc_samples["noise"] = torch.rand(num_samples, 1, **tkwargs)
@@ -190,7 +190,7 @@ def test_fit_model(self):
                 model, warmup_steps=8, num_samples=5, thinning=2, disable_progbar=True
             )
             self.assertIsInstance(model.mean_module, ConstantMean)
-            self.assertEqual(model.mean_module.constant.shape, torch.Size([3, 1]))
+            self.assertEqual(model.mean_module.raw_constant.shape, torch.Size([3]))
             self.assertIsInstance(model.covar_module, ScaleKernel)
             self.assertEqual(model.covar_module.outputscale.shape, torch.Size([3]))
             self.assertIsInstance(model.covar_module.base_kernel, MaternKernel)
@@ -465,7 +465,7 @@ def test_load_samples(self):
             )
             self.assertTrue(
                 torch.allclose(
-                    model.mean_module.constant.data,
+                    model.mean_module.raw_constant.data,
                     mcmc_samples["mean"],
                 )
             )

test/optim/test_numpy_converter.py

@@ -43,7 +43,7 @@ def test_basic(self):
             expected_sizes = {
                 "likelihood.noise_covar.raw_noise": torch.Size([1]),
                 "model.covar_module.raw_lengthscale": torch.Size([1, 3]),
-                "model.mean_module.constant": torch.Size([1]),
+                "model.mean_module.raw_constant": torch.Size(),
             }
             self.assertEqual(set(pdict.keys()), set(expected_sizes.keys()))
             for pname, val in pdict.items():
@@ -65,7 +65,7 @@ def test_exclude(self):
             mll = ExactMarginalLogLikelihood(likelihood, model)
             # test the basic case
             x, pdict, bounds = module_to_array(
-                module=mll, exclude={"model.mean_module.constant"}
+                module=mll, exclude={"model.mean_module.raw_constant"}
             )
             self.assertTrue(np.array_equal(x, np.zeros(4)))
             expected_sizes = {
@@ -98,15 +98,18 @@ def test_manual_bounds(self):
             expected_sizes = {
                 "likelihood.noise_covar.raw_noise": torch.Size([1]),
                 "model.covar_module.raw_lengthscale": torch.Size([1, 3]),
-                "model.mean_module.constant": torch.Size([1]),
+                "model.mean_module.raw_constant": torch.Size(),
             }
             self.assertEqual(set(pdict.keys()), set(expected_sizes.keys()))
             for pname, val in pdict.items():
                 self.assertEqual(val.dtype, dtype)
                 self.assertEqual(val.shape, expected_sizes[pname])
                 self.assertEqual(val.device.type, self.device.type)
             lower_exp = np.full_like(x, 0.1)
-            for p in ("likelihood.noise_covar.raw_noise", "model.mean_module.constant"):
+            for p in (
+                "likelihood.noise_covar.raw_noise",
+                "model.mean_module.raw_constant",
+            ):
                 lower_exp[_get_index(pdict, p)] = -np.inf
             self.assertTrue(np.equal(bounds[0], lower_exp).all())
             self.assertTrue(np.equal(bounds[1], np.full_like(x, np.inf)).all())
@@ -132,15 +135,15 @@ def test_module_bounds(self):
             expected_sizes = {
                 "likelihood.noise_covar.raw_noise": torch.Size([1]),
                 "model.covar_module.raw_lengthscale": torch.Size([1, 3]),
-                "model.mean_module.constant": torch.Size([1]),
+                "model.mean_module.raw_constant": torch.Size(),
             }
             self.assertEqual(set(pdict.keys()), set(expected_sizes.keys()))
             for pname, val in pdict.items():
                 self.assertEqual(val.dtype, dtype)
                 self.assertEqual(val.shape, expected_sizes[pname])
                 self.assertEqual(val.device.type, self.device.type)
             lower_exp = np.full_like(x, 0.1)
-            lower_exp[_get_index(pdict, "model.mean_module.constant")] = -np.inf
+            lower_exp[_get_index(pdict, "model.mean_module.raw_constant")] = -np.inf
             lower_exp[_get_index(pdict, "likelihood.noise_covar.raw_noise")] = 1e-5
             self.assertTrue(np.allclose(bounds[0], lower_exp))
             self.assertTrue(np.equal(bounds[1], np.full_like(x, np.inf)).all())
@@ -178,8 +181,8 @@ def test_set_parameters(self):
             )
             self.assertTrue(
                 torch.equal(
-                    z["model.mean_module.constant"],
-                    torch.tensor([5.0], device=self.device, dtype=dtype),
+                    z["model.mean_module.raw_constant"],
+                    torch.tensor(5.0, device=self.device, dtype=dtype),
                 )
             )
 

test/test_fit.py

@@ -105,7 +105,7 @@ def test_fit_gpytorch_model(self, optimizer=fit_gpytorch_scipy):
             model = mll.model
             # Make sure all of the parameters changed
             self.assertGreater(model.likelihood.raw_noise.abs().item(), 1e-3)
-            self.assertLess(model.mean_module.constant.abs().item(), 0.1)
+            self.assertLess(model.mean_module.raw_constant.abs().item(), 0.1)
             self.assertGreater(
                 model.covar_module.base_kernel.raw_lengthscale.abs().item(), 0.1
             )
@@ -123,7 +123,7 @@ def test_fit_gpytorch_model(self, optimizer=fit_gpytorch_scipy):
 
             model = mll.model
             self.assertGreaterEqual(model.likelihood.raw_noise.abs().item(), 1e-1)
-            self.assertLess(model.mean_module.constant.abs().item(), 0.1)
+            self.assertLess(model.mean_module.raw_constant.abs().item(), 0.1)
             self.assertGreater(
                 model.covar_module.base_kernel.raw_lengthscale.abs().item(), 0.1
             )
@@ -176,10 +176,10 @@ def test_fit_gpytorch_model(self, optimizer=fit_gpytorch_scipy):
 
             # test excluding a parameter
             mll = self._getModel(double=double)
-            original_raw_noise = mll.model.likelihood.noise_covar.raw_noise.item()
-            original_mean_module_constant = mll.model.mean_module.constant.item()
+            orig_raw_noise = mll.model.likelihood.noise_covar.raw_noise.item()
+            orig_mean_module_raw_constant = mll.model.mean_module.raw_constant.item()
             options["exclude"] = [
-                "model.mean_module.constant",
+                "model.mean_module.raw_constant",
                 "likelihood.noise_covar.raw_noise",
             ]
             mll = fit_gpytorch_model(
@@ -188,10 +188,10 @@ def test_fit_gpytorch_model(self, optimizer=fit_gpytorch_scipy):
             model = mll.model
             # Make excluded params did not change
             self.assertEqual(
-                model.likelihood.noise_covar.raw_noise.item(), original_raw_noise
+                model.likelihood.noise_covar.raw_noise.item(), orig_raw_noise
             )
             self.assertEqual(
-                model.mean_module.constant.item(), original_mean_module_constant
+                model.mean_module.raw_constant.item(), orig_mean_module_raw_constant
             )
             # Make sure other params did change
             self.assertGreater(
@@ -212,7 +212,7 @@ def test_fit_gpytorch_model(self, optimizer=fit_gpytorch_scipy):
             model = mll.model
             # Make sure all of the parameters changed
             self.assertGreater(model.likelihood.raw_noise.abs().item(), 1e-3)
-            self.assertLess(model.mean_module.constant.abs().item(), 0.1)
+            self.assertLess(model.mean_module.raw_constant.abs().item(), 0.1)
             self.assertGreater(
                 model.covar_module.base_kernel.raw_lengthscale.abs().item(), 0.1
             )

test/utils/test_gp_sampling.py

@@ -74,7 +74,7 @@ def _get_model(
             1.1000, **tkwargs
         ),
         "likelihood.noise_covar.noise_prior.rate": torch.tensor(0.0500, **tkwargs),
-        "mean_module.constant": torch.tensor([0.1398], **tkwargs),
+        "mean_module.raw_constant": torch.tensor(0.1398, **tkwargs),
         "covar_module.raw_outputscale": torch.tensor(0.6933, **tkwargs),
         "covar_module.base_kernel.raw_lengthscale": torch.tensor(
             [[-0.0444]], **tkwargs
@@ -111,8 +111,8 @@ def _get_model(
                 torch.tensor([0.0745], **tkwargs),
             ]
         )
-        state_dict["mean_module.constant"] = torch.stack(
-            [state_dict["mean_module.constant"], torch.tensor([0.3276], **tkwargs)]
+        state_dict["mean_module.raw_constant"] = torch.stack(
+            [state_dict["mean_module.raw_constant"], torch.tensor(0.3276, **tkwargs)]
         )
         state_dict["covar_module.raw_outputscale"] = torch.stack(
             [
@@ -134,7 +134,7 @@ def _get_model(
         state_dict["likelihood.noise_covar.raw_noise"] = torch.tensor(
             [[0.0214], [0.001]], **tkwargs
         )
-        state_dict["mean_module.constant"] = torch.tensor([[0.1398], [0.5]], **tkwargs)
+        state_dict["mean_module.raw_constant"] = torch.tensor([0.1398, 0.5], **tkwargs)
         state_dict["covar_module.raw_outputscale"] = torch.tensor(
             [0.6933, 1.0], **tkwargs
         )
@@ -153,8 +153,8 @@ def _get_model(
         state_dict["likelihood.noise_covar.raw_noise"] = torch.tensor(
             [[0.1743], [0.3132]] if multi_output else [0.1743], **tkwargs
         )
-        state_dict["mean_module.constant"] = torch.tensor(
-            [[0.2560], [0.6714]] if multi_output else [0.2555], **tkwargs
+        state_dict["mean_module.raw_constant"] = torch.tensor(
+            [0.2560, 0.6714] if multi_output else 0.2555, **tkwargs
         )
         state_dict["covar_module.raw_outputscale"] = torch.tensor(
             [2.4396, 2.6821] if multi_output else 2.4398, **tkwargs
@@ -187,15 +187,15 @@ def test_gp_draw_single_output(self):
         for dtype in (torch.float, torch.double):
             tkwargs = {"device": self.device, "dtype": dtype}
             model, _, _ = _get_model(**tkwargs)
-            mean = model.mean_module.constant.detach().clone()
+            mean = model.mean_module.raw_constant.detach().clone()
             gp = GPDraw(model)
             # test initialization
             self.assertIsNone(gp.Xs)
             self.assertIsNone(gp.Ys)
             self.assertIsNotNone(gp._seed)
             # make sure model is actually deepcopied
-            model.mean_module.constant = None
-            self.assertTrue(torch.equal(gp._model.mean_module.constant, mean))
+            model.mean_module.constant = float("inf")
+            self.assertTrue(torch.equal(gp._model.mean_module.raw_constant, mean))
             # test basic functionality
             test_X1 = torch.rand(1, 1, **tkwargs, requires_grad=True)
             Y1 = gp(test_X1)
@@ -234,14 +234,14 @@ def test_gp_draw_multi_output(self):
         for dtype in (torch.float, torch.double):
             tkwargs = {"device": self.device, "dtype": dtype}
             model, _, _ = _get_model(**tkwargs, multi_output=True)
-            mean = model.mean_module.constant.detach().clone()
+            mean = model.mean_module.raw_constant.detach().clone()
             gp = GPDraw(model)
             # test initialization
             self.assertIsNone(gp.Xs)
             self.assertIsNone(gp.Ys)
             # make sure model is actually deepcopied
-            model.mean_module.constant = None
-            self.assertTrue(torch.equal(gp._model.mean_module.constant, mean))
+            model.mean_module.constant = float("inf")
+            self.assertTrue(torch.equal(gp._model.mean_module.raw_constant, mean))
             # test basic functionality
             test_X1 = torch.rand(1, 1, **tkwargs, requires_grad=True)
             Y1 = gp(test_X1)