modify PenalizedMCObjective to support non-batch eval (#2073)

Summary:
Pull Request resolved: #2073

To match penalty term with MC objective, we current unsqueeze the first dim which corresponds to the dimension of MC samples. However, when a `qxd`-dim X tensor is evaluated e.g. computing feasibility, it causes shape mismatch. As one would expect `q`-dim tensor returned, it will return `1xq`-dim tensor instead.

To fix, we check the dims of obj; if it is non-mc samples, we will sequeeze the first dim back.

Reviewed By: bletham

Differential Revision: D49305807

fbshipit-source-id: c8e2dedbdb38d79b5910aca393cf2d3d4e1e311f

Author: Qing Feng

botorch/acquisition/penalized.py

@@ -366,6 +366,12 @@ def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
         if self.expand_dim is not None:
             # reshape penalty_obj to match the dim
             penalty_obj = penalty_obj.unsqueeze(self.expand_dim)
+        # this happens when samples is a `q x m`-dim tensor and X is a `q x d`-dim
+        # tensor; obj returned from GenericMCObjective is a `q`-dim tensor and
+        # penalty_obj is a `1 x q`-dim tensor.
+        if obj.ndim == 1:
+            assert penalty_obj.shape == torch.Size([1, samples.shape[-2]])
+            penalty_obj = penalty_obj.squeeze(dim=0)
         return obj - self.regularization_parameter * penalty_obj
 
 

test/acquisition/test_penalized.py

@@ -291,7 +291,7 @@ def test_penalized_mc_objective(self):
             samples = torch.randn(4, 3, device=self.device, dtype=dtype)
             X = torch.randn(4, 5, device=self.device, dtype=dtype)
             penalized_obj = generic_obj(samples) - 0.1 * l1_penalty_obj(X)
-            self.assertTrue(torch.equal(obj(samples, X), penalized_obj))
+            self.assertTrue(torch.equal(obj(samples, X), penalized_obj.squeeze(0)))
             # test 'q x d' Tensor X
             samples = torch.randn(4, 2, 3, device=self.device, dtype=dtype)
             X = torch.randn(2, 5, device=self.device, dtype=dtype)