Fix shape error in qNegIntegratedPosteriorVariance (#3068)

Summary:
Pull Request resolved: #3068

`qNegIntegratedPosteriorVariance` doesn't work through MBM due to a shape bug. This diff addresses the underlying issue and would allow us to use active learning through Ax (at some point).

Reviewed By: saitcakmak

Differential Revision: D85507279

fbshipit-source-id: 1917902e3abd989596147bd4bb74f934697cb5df

Author: David Eriksson

botorch/acquisition/active_learning.py

@@ -120,13 +120,12 @@ def forward(self, X: Tensor) -> Tensor:
             )
 
         neg_variance = posterior.variance.mul(-1.0)
-
-        if self.posterior_transform is None:
-            # if single-output, shape is 1 x batch_shape x num_grid_points x 1
-            return neg_variance.mean(dim=-2).squeeze(-1).squeeze(0)
-        else:
+        if self.model.num_outputs > 1:
             # if multi-output + obj, shape is num_grid_points x batch_shape x 1 x 1
             return neg_variance.mean(dim=0).squeeze(-1).squeeze(-1)
+        else:
+            # if single-output, shape is 1 x batch_shape x num_grid_points x 1
+            return neg_variance.mean(dim=-2).squeeze(-1).squeeze(0)
 
 
 class PairwiseMCPosteriorVariance(MCAcquisitionFunction):

test/acquisition/test_active_learning.py

@@ -132,6 +132,29 @@ def test_q_neg_int_post_variance(self):
                     val_exp = -0.5 * variance.mean(dim=0).view(3, -1).mean(dim=-1)
                     self.assertAllClose(val, val_exp, atol=1e-4)
 
+            # single-output with posterior transform
+            mean = torch.zeros(2, 10, 3, 1, device=self.device, dtype=dtype)
+            variance = torch.ones(2, 10, 3, 1, device=self.device, dtype=dtype)
+            cov = torch.diag_embed(variance.view(2, 10, -1))
+            f_posterior = GPyTorchPosterior(MultitaskMultivariateNormal(mean, cov))
+            mc_points = torch.rand(10, 1, device=self.device, dtype=dtype)
+            mfm = MockModel(f_posterior)
+            with mock.patch.object(MockModel, "fantasize", return_value=mfm):
+                with mock.patch(no, new_callable=mock.PropertyMock) as mock_num_outputs:
+                    mock_num_outputs.return_value = 1
+                    mm = MockModel(None)
+                    qNIPV = qNegIntegratedPosteriorVariance(
+                        model=mm,
+                        mc_points=mc_points,
+                        posterior_transform=ScalarizedPosteriorTransform(
+                            weights=torch.tensor([1.0], device=self.device, dtype=dtype)
+                        ),
+                    )
+                    X = torch.empty(2, 3, 1, 1, device=self.device, dtype=dtype)
+                    val = qNIPV(X)
+                    val_exp = -variance.mean(dim=-2).squeeze(-1).squeeze(0)
+                    self.assertAllClose(val, val_exp, atol=1e-4)
+
 
 class TestPairwiseMCPosteriorVariance(BotorchTestCase):
     def test_pairwise_mc_post_var(self):