Fix unit tests broken by upstream gpytorch changes

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

Author: Balandat

ax/models/tests/test_alebo.py

@@ -89,7 +89,7 @@ def testALEBOGP(self):
 
         # Batch
         Uvec_b = m.covar_module.base_kernel.Uvec.repeat(5, 1)
-        mean_b = m.mean_module.constant.repeat(5, 1)
+        mean_b = m.mean_module.constant.repeat(5)
         output_scale_b = m.covar_module.raw_outputscale.repeat(5)
         m_b = get_batch_model(
             B=B,
@@ -132,7 +132,7 @@ def testALEBOGP(self):
             {
                 "covar_module.base_kernel.Uvec",
                 "covar_module.raw_outputscale",
-                "mean_module.constant",
+                "mean_module.raw_constant",
                 "covar_module.raw_outputscale_constraint.lower_bound",
                 "covar_module.raw_outputscale_constraint.upper_bound",
             },
@@ -151,7 +151,7 @@ def testALEBOGP(self):
                 {
                     "covar_module.base_kernel.Uvec",
                     "covar_module.raw_outputscale",
-                    "mean_module.constant",
+                    "mean_module.raw_constant",
                     "covar_module.raw_outputscale_constraint.lower_bound",
                     "covar_module.raw_outputscale_constraint.upper_bound",
                 },

ax/models/tests/test_botorch_model.py

@@ -457,7 +457,7 @@ def test_BotorchModel(self, dtype=torch.float, cuda=False):
 
             # Test loading state dict
             true_state_dict = {
-                "mean_module.constant": [3.5004],
+                "mean_module.raw_constant": 3.5004,
                 "covar_module.raw_outputscale": 2.2438,
                 "covar_module.base_kernel.raw_lengthscale": [
                     [-0.9274, -0.9274, -0.9274]
@@ -489,7 +489,7 @@ def test_BotorchModel(self, dtype=torch.float, cuda=False):
                 self.assertTrue(torch.equal(true_state_dict[k], v))
 
             # Test for some change in model parameters & buffer for refit_model=True
-            true_state_dict["mean_module.constant"] += 0.1
+            true_state_dict["mean_module.raw_constant"] += 0.1
             true_state_dict["covar_module.raw_outputscale"] += 0.1
             true_state_dict["covar_module.base_kernel.raw_lengthscale"] += 0.1
             model = get_and_fit_model(

ax/models/tests/test_fully_bayesian.py

@@ -348,7 +348,7 @@ def test_FullyBayesianBotorchModel(self, dtype=torch.float, cuda=False):
                 self.assertEqual(len(model.model.models), 2)
                 m1, m2 = model.model.models[0], model.model.models[1]
                 # Mean
-                self.assertEqual(m1.mean_module.constant.shape, (4, 1))
+                self.assertEqual(m1.mean_module.constant.shape, (4,))
                 self.assertFalse(
                     torch.isclose(
                         m1.mean_module.constant, m2.mean_module.constant
@@ -376,11 +376,9 @@ def test_FullyBayesianBotorchModel(self, dtype=torch.float, cuda=False):
                 device = torch.device("cuda") if cuda else torch.device("cpu")
                 objective_weights = torch.tensor([1.0, 0.0], dtype=dtype, device=device)
                 objective_transform = get_objective_weights_transform(objective_weights)
-                infeasible_cost = torch.tensor(
-                    get_infeasible_cost(
+                infeasible_cost = get_infeasible_cost(
                         X=Xs1[0], model=model.model, objective=objective_transform
-                    )
-                )
+                    ).detach().clone()
                 expected_infeasible_cost = -1 * torch.min(
                     objective_transform(
                         model.model.posterior(Xs1[0]).mean
@@ -841,7 +839,7 @@ def test_FullyBayesianBotorchModelPyro(self, dtype=torch.double, cuda=False):
                         )
                         self.assertEqual(
                             m.mean_module.constant.shape,
-                            torch.Size([4, 1]),
+                            torch.Size([4]),
                         )
                         if use_input_warping:
                             self.assertTrue(hasattr(m, "input_transform"))

ax/models/torch/alebo.py

@@ -331,7 +331,7 @@ def f(x):
 
     # Sample only Uvec; leave mean and output scale fixed.
     assert list(property_dict.keys()) == [
-        "model.mean_module.constant",
+        "model.mean_module.raw_constant",
         "model.covar_module.raw_outputscale",
         "model.covar_module.base_kernel.Uvec",
     ]
@@ -347,7 +347,7 @@ def f(x):
         nsamp, *attrs.shape
     )
     # Get the other properties into batch mode
-    mean_constant_batch = mll.model.mean_module.constant.repeat(nsamp, 1)
+    mean_constant_batch = mll.model.mean_module.constant.repeat(nsamp)
     output_scale_batch = mll.model.covar_module.raw_outputscale.repeat(nsamp)
     return Uvec_batch, mean_constant_batch, output_scale_batch
 
@@ -383,10 +383,10 @@ def get_batch_model(
     )
     m_b.train()
     # Set mean constant
-    # pyre-fixme[16]: `Optional` has no attribute `constant`.
-    m_b.mean_module.constant.requires_grad_(False)
-    m_b.mean_module.constant.copy_(mean_constant_batch)
-    m_b.mean_module.constant.requires_grad_(True)
+    # pyre-fixme[16]: `Optional` has no attribute `raw_constant`.
+    m_b.mean_module.raw_constant.requires_grad_(False)
+    m_b.mean_module.raw_constant.copy_(mean_constant_batch)
+    m_b.mean_module.raw_constant.requires_grad_(True)
     # Set output scale
     m_b.covar_module.raw_outputscale.requires_grad_(False)
     m_b.covar_module.raw_outputscale.copy_(output_scale_batch)