Fix inconsistent output shapes when all features are fixed in optimize_acqf (#3241)

Summary:
When all features are fixed, _optimize_acqf_all_features_fixed now returns shapes consistent with the normal _optimize_acqf_batch path:
- return_best_only=True: candidates (q, d), acq_value scalar
- return_best_only=False: candidates (1, q, d), acq_value (1,)

Closes #2740

Pull Request resolved: #3241

Reviewed By: dme65

Differential Revision: D97315513

Pulled By: saitcakmak

fbshipit-source-id: 28a44486f72706ccdf77d3a77faa5cce6ec9ef69

Author: saitcakmak

botorch/optim/optimize.py

@@ -198,6 +198,7 @@ def _optimize_acqf_all_features_fixed(
     fixed_features: dict[int, float],
     q: int,
     acq_function: AcquisitionFunction,
+    return_best_only: bool = True,
     return_acq_values: bool = True,
 ) -> tuple[Tensor, Tensor | None]:
     """
@@ -210,10 +211,21 @@ def _optimize_acqf_all_features_fixed(
         dtype=bounds.dtype,
     )
     X = X.expand(q, *X.shape)
+    if not return_best_only:
+        # When return_best_only=False, candidates have shape
+        # `num_restarts x q x d`. With all features fixed there is only one
+        # candidate, so num_restarts=1.
+        X = X.unsqueeze(0)
     if not return_acq_values:
         return X, None
     with torch.no_grad():
-        acq_value = acq_function(X)
+        acq_value = acq_function(X if return_best_only else X.squeeze(0))
+    # Ensure acq_value is a scalar (consistent with return_best_only=True)
+    # or 1-d with shape `(1,)` (consistent with return_best_only=False).
+    if return_best_only:
+        acq_value = acq_value.squeeze()
+    else:
+        acq_value = acq_value.view(1)
     return X, acq_value
 
 
@@ -810,6 +822,7 @@ def _optimize_acqf(opt_inputs: OptimizeAcqfInputs) -> tuple[Tensor, Tensor | Non
             fixed_features=opt_inputs.fixed_features,
             q=opt_inputs.q,
             acq_function=opt_inputs.acq_function,
+            return_best_only=opt_inputs.return_best_only,
             return_acq_values=opt_inputs.return_acq_values,
         )
 

test/optim/test_optimize.py

@@ -407,6 +407,36 @@ def test_optimize_acqf_return_acq_values(
             )
         )
 
+        # All features fixed: return shape consistency with normal path.
+        # return_best_only=True (default): candidates (q, d), acq_value scalar
+        candidates_rbo, acq_rbo = optimize_acqf(
+            acq_function=mock_acq_function,
+            bounds=bounds,
+            q=1,
+            num_restarts=num_restarts,
+            raw_samples=raw_samples,
+            options=options,
+            fixed_features=fixed_all,
+            return_best_only=True,
+        )
+        self.assertEqual(candidates_rbo.shape, (1, 3))
+        self.assertEqual(acq_rbo.shape, torch.Size([]))
+
+        # return_best_only=False: candidates (num_restarts, q, d),
+        # acq_value (num_restarts,)
+        candidates_all, acq_all = optimize_acqf(
+            acq_function=mock_acq_function,
+            bounds=bounds,
+            q=1,
+            num_restarts=num_restarts,
+            raw_samples=raw_samples,
+            options=options,
+            fixed_features=fixed_all,
+            return_best_only=False,
+        )
+        self.assertEqual(candidates_all.shape, (1, 1, 3))
+        self.assertEqual(acq_all.shape, (1,))
+
         # Sequential path: return_acq_values=True and return_acq_values=False
         mock_gen_candidates_scipy.return_value = (
             torch.rand(1, 1, 3, device=self.device, dtype=torch.double),