fix: Introduce parameter to other optimizers for n batch

automl · Jan 8, 2025 · d41daac · d41daac
1 parent 1d646ed
commit d41daac
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 2 deletions.
diff --git a/neps/optimizers/grid_search/optimizer.py b/neps/optimizers/grid_search/optimizer.py
@@ -100,7 +100,9 @@ def ask(
         self,
         trials: Mapping[str, Trial],
         budget_info: BudgetInfo | None,
+        n: int | None = None,
     ) -> SampledConfig:
+        assert n is None, "TODO"
         _num_previous_configs = len(trials)
         if _num_previous_configs > len(self.configs_list) - 1:
             raise ValueError("Grid search exhausted!")

diff --git a/neps/optimizers/multi_fidelity/hyperband.py b/neps/optimizers/multi_fidelity/hyperband.py
@@ -126,7 +126,9 @@ def ask(
         self,
         trials: Mapping[str, Trial],
         budget_info: BudgetInfo | None,
+        n: int | None = None,
     ) -> SampledConfig:
+        assert n is None, "TODO"
         completed: dict[str, ConfigResult] = {
             trial_id: trial.into_config_result(self.pipeline_space.from_dict)
             for trial_id, trial in trials.items()

diff --git a/neps/optimizers/multi_fidelity/ifbo.py b/neps/optimizers/multi_fidelity/ifbo.py
@@ -171,8 +171,10 @@ def __init__(
     def ask(
         self,
         trials: Mapping[str, Trial],
-        max_cost_total_info: BudgetInfo | None = None,
+        budget_info: BudgetInfo | None = None,
+        n: int | None = None,
     ) -> SampledConfig:
+        assert n is None, "TODO"
         ids = [int(config_id.split("_", maxsplit=1)[0]) for config_id in trials]
         new_id = max(ids) + 1 if len(ids) > 0 else 0
 

diff --git a/neps/optimizers/multi_fidelity/successive_halving.py b/neps/optimizers/multi_fidelity/successive_halving.py
@@ -321,8 +321,10 @@ def ask(
         self,
         trials: Mapping[str, Trial],
         budget_info: BudgetInfo | None,
+        n: int | None = None,
     ) -> SampledConfig:
         """This is basically the fit method."""
+        assert n is None, "TODO"
         completed: dict[str, ConfigResult] = {
             trial_id: trial.into_config_result(self.pipeline_space.from_dict)
             for trial_id, trial in trials.items()

diff --git a/neps/optimizers/multi_fidelity_prior/async_priorband.py b/neps/optimizers/multi_fidelity_prior/async_priorband.py
@@ -238,9 +238,11 @@ def _update_sh_bracket_state(self) -> None:
     def ask(
         self,
         trials: Mapping[str, Trial],
-        max_cost_total_info: BudgetInfo | None,
+        budget_info: BudgetInfo | None,
+        n: int | None = None,
     ) -> SampledConfig:
         """This is basically the fit method."""
+        assert n is None, "TODO"
         completed: dict[str, ConfigResult] = {
             trial_id: trial.into_config_result(self.pipeline_space.from_dict)
             for trial_id, trial in trials.items()