Validate output shape of models upon instantiating acq. functions (#331)

Summary:
Pull Request resolved: #331

Addresses #328

Reviewed By: liangshi7

Differential Revision: D18739840

fbshipit-source-id: 4b888d17a24de211315d36a09ce4e879d294973e

Author: Balandat

botorch/acquisition/analytic.py

@@ -41,37 +41,36 @@ def __init__(
             objective: A ScalarizedObjective (optional).
         """
         super().__init__(model=model)
-        if objective is not None and not isinstance(objective, ScalarizedObjective):
+        if objective is None:
+            if model.num_outputs != 1:
+                raise UnsupportedError(
+                    "Must specify an objective when using a multi-output model."
+                )
+        elif not isinstance(objective, ScalarizedObjective):
             raise UnsupportedError(
                 "Only objectives of type ScalarizedObjective are supported for "
                 "analytic acquisition functions."
             )
         self.objective = objective
 
-    def _get_posterior(self, X: Tensor, check_single_output: bool = True) -> Posterior:
+    def _get_posterior(self, X: Tensor) -> Posterior:
         r"""Compute the posterior at the input candidate set X.
 
         Applies the objective if provided.
 
         Args:
-            X: The input candidate set
-            check_single_output: If True, Raise an error if the posterior is not
-                single-output.
+            X: The input candidate set.
 
         Returns:
-            The posterior at X.
+            The posterior at X. If a ScalarizedObjective is defined, this
+            posterior can be single-output even if the underlying model is a
+            multi-output model.
         """
         posterior = self.model.posterior(X)
         if self.objective is not None:
-            # Unlike MCAcquisitionObjectives (which transform samples), this
+            # Unlike MCAcquisitionObjective (which transform samples), this
             # transforms the posterior
             posterior = self.objective(posterior)
-        if check_single_output:
-            if posterior.event_shape[-1] != 1:
-                raise UnsupportedError(
-                    "Multi-Output posteriors are not supported for acquisition "
-                    f"functions of type {self.__class__.__name__}"
-                )
         return posterior
 
     def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
@@ -349,7 +348,9 @@ def __init__(
                 bounds on that output (resp. interpreted as -Inf / Inf if None)
             maximize: If True, consider the problem a maximization problem.
         """
-        super().__init__(model=model)
+        # use AcquisitionFunction constructor to avoid check for objective
+        super(AnalyticAcquisitionFunction, self).__init__(model=model)
+        self.objective = None
         self.maximize = maximize
         self.objective_index = objective_index
         self.constraints = constraints
@@ -369,7 +370,7 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b)`-dim Tensor of Expected Improvement values at the given
             design points `X`.
         """
-        posterior = self._get_posterior(X=X, check_single_output=False)
+        posterior = self._get_posterior(X=X)
         means = posterior.mean.squeeze(dim=-2)  # (b) x m
         sigmas = posterior.variance.squeeze(dim=-2).sqrt().clamp_min(1e-9)  # (b) x m
 

botorch/acquisition/max_value_entropy_search.py

@@ -89,13 +89,6 @@ def __init__(
         sampler = SobolQMCNormalSampler(num_y_samples)
         super().__init__(model=model, sampler=sampler)
 
-        # The following check is specific to batched models
-        # @TODO: model-agnostic num_outputs check: #295
-        if self.model._num_outputs > 1:
-            raise NotImplementedError(
-                "Models with > 1 outcomes are not yet supported by qMaxValueEntropy!"
-            )
-
         # Batch GP models (e.g. fantasized models) are not currently supported
         if self.model.train_inputs[0].ndim > 2:
             raise NotImplementedError(

botorch/acquisition/monte_carlo.py

@@ -65,6 +65,10 @@ def __init__(
             sampler = SobolQMCNormalSampler(num_samples=512, collapse_batch_dims=True)
         self.add_module("sampler", sampler)
         if objective is None:
+            if model.num_outputs != 1:
+                raise UnsupportedError(
+                    "Must specify an objective when using a multi-output model."
+                )
             objective = IdentityMCObjective()
         elif not isinstance(objective, MCAcquisitionObjective):
             raise UnsupportedError(

test/acquisition/test_analytic.py

@@ -38,10 +38,21 @@
 ]
 
 
+class DummyAnalyticAcquisitionFunction(AnalyticAcquisitionFunction):
+    def forward(self, X):
+        pass
+
+
 class TestAnalyticAcquisitionFunction(BotorchTestCase):
     def test_abstract_raises(self):
         with self.assertRaises(TypeError):
             AnalyticAcquisitionFunction()
+        # raise if model is multi-output, but no objective is given
+        mean = torch.zeros(1, 2)
+        variance = torch.ones(1, 2)
+        mm = MockModel(MockPosterior(mean=mean, variance=variance))
+        with self.assertRaises(UnsupportedError):
+            DummyAnalyticAcquisitionFunction(model=mm)
 
 
 class TestExpectedImprovement(BotorchTestCase):
@@ -113,9 +124,8 @@ def test_expected_improvement_batch(self):
             mean2 = torch.rand(3, 1, 2, device=self.device, dtype=dtype)
             variance2 = torch.rand(3, 1, 2, device=self.device, dtype=dtype)
             mm2 = MockModel(MockPosterior(mean=mean2, variance=variance2))
-            module2 = ExpectedImprovement(model=mm2, best_f=0.0)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                ExpectedImprovement(model=mm2, best_f=0.0)
 
             # test objective (single-output)
             mean = torch.tensor([[[0.5]], [[0.25]]], device=self.device, dtype=dtype)
@@ -172,9 +182,8 @@ def test_posterior_mean(self):
             # check for proper error if multi-output model
             mean2 = torch.rand(1, 2, device=self.device, dtype=dtype)
             mm2 = MockModel(MockPosterior(mean=mean2))
-            module2 = PosteriorMean(model=mm2)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                PosteriorMean(model=mm2)
 
     def test_posterior_mean_batch(self):
         for dtype in (torch.float, torch.double):
@@ -189,9 +198,8 @@ def test_posterior_mean_batch(self):
             # check for proper error if multi-output model
             mean2 = torch.rand(3, 1, 2, device=self.device, dtype=dtype)
             mm2 = MockModel(MockPosterior(mean=mean2))
-            module2 = PosteriorMean(model=mm2)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                PosteriorMean(model=mm2)
 
 
 class TestProbabilityOfImprovement(BotorchTestCase):
@@ -217,9 +225,8 @@ def test_probability_of_improvement(self):
             mean2 = torch.rand(1, 2, device=self.device, dtype=dtype)
             variance2 = torch.ones_like(mean2)
             mm2 = MockModel(MockPosterior(mean=mean2, variance=variance2))
-            module2 = ProbabilityOfImprovement(model=mm2, best_f=0.0)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                ProbabilityOfImprovement(model=mm2, best_f=0.0)
 
     def test_probability_of_improvement_batch(self):
         for dtype in (torch.float, torch.double):
@@ -237,9 +244,8 @@ def test_probability_of_improvement_batch(self):
             mean2 = torch.rand(3, 1, 2, device=self.device, dtype=dtype)
             variance2 = torch.ones_like(mean2)
             mm2 = MockModel(MockPosterior(mean=mean2, variance=variance2))
-            module2 = ProbabilityOfImprovement(model=mm2, best_f=0.0)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                ProbabilityOfImprovement(model=mm2, best_f=0.0)
 
 
 class TestUpperConfidenceBound(BotorchTestCase):
@@ -265,9 +271,8 @@ def test_upper_confidence_bound(self):
             mean2 = torch.rand(1, 2, device=self.device, dtype=dtype)
             variance2 = torch.rand(1, 2, device=self.device, dtype=dtype)
             mm2 = MockModel(MockPosterior(mean=mean2, variance=variance2))
-            module2 = UpperConfidenceBound(model=mm2, beta=1.0)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                UpperConfidenceBound(model=mm2, beta=1.0)
 
     def test_upper_confidence_bound_batch(self):
         for dtype in (torch.float, torch.double):
@@ -287,9 +292,8 @@ def test_upper_confidence_bound_batch(self):
             mean2 = torch.rand(3, 1, 2, device=self.device, dtype=dtype)
             variance2 = torch.rand(3, 1, 2, device=self.device, dtype=dtype)
             mm2 = MockModel(MockPosterior(mean=mean2, variance=variance2))
-            module2 = UpperConfidenceBound(model=mm2, beta=1.0)
             with self.assertRaises(UnsupportedError):
-                module2(X)
+                UpperConfidenceBound(model=mm2, beta=1.0)
 
 
 class TestConstrainedExpectedImprovement(BotorchTestCase):