add branin50, hartmann50 to benchmark registry

ax/benchmark/problems/registry.py

@@ -77,6 +77,16 @@ class BenchmarkProblemRegistryEntry:
             "observe_noise_sd": False,
         },
     ),
+    "branin50": BenchmarkProblemRegistryEntry(
+        factory_fn=create_problem_from_botorch,
+        factory_kwargs={
+            "test_problem_class": synthetic.Branin,
+            "test_problem_kwargs": {},
+            "num_trials": 50,
+            "observe_noise_sd": False,
+            "n_dummy_dimensions": 48,
+        },
+    ),
     "branin_currin": BenchmarkProblemRegistryEntry(
         factory_fn=create_problem_from_botorch,
         factory_kwargs={
@@ -134,6 +144,16 @@ class BenchmarkProblemRegistryEntry:
             "n_dummy_dimensions": 24,
         },
     ),
+    "hartmann50": BenchmarkProblemRegistryEntry(
+        factory_fn=create_problem_from_botorch,
+        factory_kwargs={
+            "test_problem_class": synthetic.Hartmann,
+            "test_problem_kwargs": {"dim": 6},
+            "num_trials": 50,
+            "observe_noise_sd": False,
+            "n_dummy_dimensions": 44,
+        },
+    ),
     "jenatton": BenchmarkProblemRegistryEntry(
         factory_fn=get_jenatton_benchmark_problem,
         factory_kwargs={"num_trials": 50, "observe_noise_sd": False},

ax/generators/torch/botorch_modular/acquisition.py

@@ -36,6 +36,7 @@
 from botorch.acquisition.acquisition import AcquisitionFunction
 from botorch.acquisition.input_constructors import get_acqf_input_constructor
 from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
+from botorch.acquisition.logei import qLogProbabilityOfFeasibility
 from botorch.acquisition.multioutput_acquisition import MultiOutputAcquisitionFunction
 from botorch.acquisition.objective import MCAcquisitionObjective, PosteriorTransform
 from botorch.exceptions.errors import BotorchError, InputDataError
@@ -54,6 +55,7 @@
 
 try:
     from botorch.utils.multi_objective.optimize import optimize_with_nsgaii
+
 except ImportError:
     optimize_with_nsgaii = None
 
@@ -670,7 +672,9 @@ def optimize(
                 f"Unknown optimizer: {optimizer}. This code should be unreachable."
             )
         # prune irrelevant parameters post-hoc
-        if self.options.get("prune_irrelevant_parameters", False):
+        if self.options.get("prune_irrelevant_parameters", False) and not isinstance(
+            self.acqf, qLogProbabilityOfFeasibility
+        ):
             if self._pruning_target_point is None:
                 logger.info(
                     "Must specify pruning_target_point to prune irrelevant "
@@ -827,7 +831,7 @@ def _prune_irrelevant_parameters(
             self._pruning_target_point,
             message="Must specify pruning_target_point to prune irrelevant parameters",
         )
-        irrelevance_pruning_rtol = self.options.get("irrelevance_pruning_rtol", 0.1)
+        irrelevance_pruning_rtol = self.options.get("irrelevance_pruning_rtol", 0.2)
         initial_X_pending = self.X_pending
         pruned_af_vals = []
         excluded_indices = set(search_space_digest.fidelity_features)

ax/generators/torch/tests/test_acquisition.py

@@ -1479,7 +1479,7 @@ def test_prune_irrelevant_parameters_no_pruning_above_threshold(self) -> None:
             side_effect=[
                 torch.tensor([0.0]),  # baseline acquisition value
                 torch.tensor([1.0]),  # original dense acquisition value
-                torch.tensor([0.8, 0.7]),
+                torch.tensor([0.7, 0.3]),
             ]
         )
         acq.evaluate = mock_evaluate
@@ -1676,6 +1676,33 @@ def test_prune_irrelevant_parameters_hss(self) -> None:
             torch.equal(pruned_values, torch.tensor([0.95], dtype=torch.double))
         )
 
+    @mock_botorch_optimize
+    def test_no_pruning_with_qLogProbabilityOfFeasibility(self) -> None:
+        # Test that pruning is NOT called when using qLogProbabilityOfFeasibility,
+        # even when prune_irrelevant_parameters option is enabled
+        self.options = {"prune_irrelevant_parameters": True}
+        self.botorch_acqf_class = qLogProbabilityOfFeasibility  # pyre-ignore [8]
+        self.botorch_acqf_options = {}
+        acquisition = self.get_acquisition_function(
+            fixed_features=self.fixed_features,
+        )
+        n = 1
+        with mock.patch.object(
+            acquisition,
+            "_prune_irrelevant_parameters",
+            wraps=acquisition._prune_irrelevant_parameters,
+        ) as mock_prune_irrelevant_parameters:
+            acquisition.optimize(
+                n=n,
+                search_space_digest=self.search_space_digest,
+                inequality_constraints=self.inequality_constraints,
+                fixed_features=self.fixed_features,
+                rounding_func=self.rounding_func,
+                optimizer_options=self.optimizer_options,
+            )
+            mock_prune_irrelevant_parameters.assert_not_called()
+            self.assertIsNone(acquisition.num_pruned_dims)
+
 
 class MultiAcquisitionTest(AcquisitionTest):
     acquisition_class = MultiAcquisition