Support customizing num_trials in high dim benchmarks (#2129)

Summary:
Pull Request resolved: #2129

Makes it possible to do `get_problem("hartmann30", num_trials=30)`.

Reviewed By: dme65

Differential Revision: D52709286

fbshipit-source-id: 5006b6dd5b42222549770a0a78300d1397a0a3e4

ax/benchmark/problems/registry.py

@@ -54,16 +54,16 @@ class BenchmarkProblemRegistryEntry:
         },
     ),
     "branin_currin30": BenchmarkProblemRegistryEntry(
-        factory_fn=lambda n: embed_higher_dimension(
+        factory_fn=lambda n, num_trials: embed_higher_dimension(
             problem=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective(
                 test_problem_class=BraninCurrin,
                 test_problem_kwargs={},
-                num_trials=30,
+                num_trials=num_trials,
                 infer_noise=True,
             ),
             total_dimensionality=n,
         ),
-        factory_kwargs={"n": 30},
+        factory_kwargs={"n": 30, "num_trials": 30},
     ),
     "griewank4": BenchmarkProblemRegistryEntry(
         factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic,
@@ -93,16 +93,16 @@ class BenchmarkProblemRegistryEntry:
         },
     ),
     "hartmann30": BenchmarkProblemRegistryEntry(
-        factory_fn=lambda n: embed_higher_dimension(
+        factory_fn=lambda n, num_trials: embed_higher_dimension(
             problem=SingleObjectiveBenchmarkProblem.from_botorch_synthetic(
                 test_problem_class=synthetic.Hartmann,
                 test_problem_kwargs={"dim": 6},
-                num_trials=25,
+                num_trials=num_trials,
                 infer_noise=True,
             ),
             total_dimensionality=n,
         ),
-        factory_kwargs={"n": 30},
+        factory_kwargs={"n": 30, "num_trials": 25},
     ),
     "hpo_pytorch_cnn_MNIST": BenchmarkProblemRegistryEntry(
         factory_fn=PyTorchCNNTorchvisionBenchmarkProblem.from_dataset_name,
@@ -181,16 +181,16 @@ class BenchmarkProblemRegistryEntry:
         },
     ),
     "branin_currin30_fixed_noise": BenchmarkProblemRegistryEntry(
-        factory_fn=lambda n: embed_higher_dimension(
+        factory_fn=lambda n, num_trials: embed_higher_dimension(
             problem=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective(
                 test_problem_class=BraninCurrin,
                 test_problem_kwargs={},
-                num_trials=100,
+                num_trials=num_trials,
                 infer_noise=False,
             ),
             total_dimensionality=n,
         ),
-        factory_kwargs={"n": 30},
+        factory_kwargs={"n": 30, "num_trials": 30},
     ),
     "hartmann6_fixed_noise": BenchmarkProblemRegistryEntry(
         factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic,
@@ -202,16 +202,16 @@ class BenchmarkProblemRegistryEntry:
         },
     ),
     "hartmann30_fixed_noise": BenchmarkProblemRegistryEntry(
-        factory_fn=lambda n: embed_higher_dimension(
+        factory_fn=lambda n, num_trials: embed_higher_dimension(
             problem=SingleObjectiveBenchmarkProblem.from_botorch_synthetic(
                 test_problem_class=synthetic.Hartmann,
                 test_problem_kwargs={"dim": 6},
-                num_trials=100,
+                num_trials=num_trials,
                 infer_noise=False,
             ),
             total_dimensionality=n,
         ),
-        factory_kwargs={"n": 30},
+        factory_kwargs={"n": 30, "num_trials": 25},
     ),
     "jenatton_fixed_noise": BenchmarkProblemRegistryEntry(
         factory_fn=get_jenatton_benchmark_problem,