Fixed bug where optimize_acqf didn't work with different batch sizes (#1668)

Summary:
Pull Request resolved: #1668

X-link: facebook/Ax#1414

Reviewed By: Balandat

Differential Revision: D43178298

fbshipit-source-id: 531dd9f62142630ea07c58c274ea6b62d48d5d2e

Author: esantorella

botorch/optim/optimize.py

@@ -342,7 +342,11 @@ def _optimize_batch_candidates(
             logger.info(f"Generated candidate batch {i+1} of {len(batched_ics)}.")
 
         batch_candidates = torch.cat(batch_candidates_list)
-        batch_acq_values = torch.stack(batch_acq_values_list).flatten()
+        has_scalars = batch_acq_values_list[0].ndim == 0
+        if has_scalars:
+            batch_acq_values = torch.stack(batch_acq_values_list)
+        else:
+            batch_acq_values = torch.cat(batch_acq_values_list).flatten()
         return batch_candidates, batch_acq_values, opt_warnings
 
     batch_candidates, batch_acq_values, ws = _optimize_batch_candidates(timeout_sec)

test/optim/test_optimize.py

@@ -382,6 +382,36 @@ def test_optimize_acqf_sequential_notimplemented(self):
                 sequential=True,
             )
 
+    def test_optimize_acqf_batch_limit(self) -> None:
+        num_restarts = 3
+        raw_samples = 5
+        dim = 4
+        q = 4
+        batch_limit = 2
+
+        options = {"batch_limit": batch_limit}
+        initial_conditions = [
+            torch.ones(shape) for shape in [(1, 2, dim), (2, 1, dim), (1, dim)]
+        ] + [None]
+
+        for gen_candidates, ics in zip(
+            [gen_candidates_scipy, gen_candidates_torch], initial_conditions
+        ):
+            with self.subTest(gen_candidates=gen_candidates, initial_conditions=ics):
+                _, acq_value_list = optimize_acqf(
+                    acq_function=SinOneOverXAcqusitionFunction(),
+                    bounds=torch.stack([-1 * torch.ones(dim), torch.ones(dim)]),
+                    q=q,
+                    num_restarts=num_restarts,
+                    raw_samples=raw_samples,
+                    options=options,
+                    return_best_only=False,
+                    gen_candidates=gen_candidates,
+                    batch_initial_conditions=ics,
+                )
+                expected_shape = (num_restarts,) if ics is None else (ics.shape[0],)
+                self.assertEqual(acq_value_list.shape, expected_shape)
+
     def test_optimize_acqf_runs_given_batch_initial_conditions(self):
         num_restarts, raw_samples, dim = 1, 2, 3