From 0b39250b8760718a2e74c9be41263d1c570f855f Mon Sep 17 00:00:00 2001 From: Elizabeth Santorella Date: Wed, 4 Sep 2024 14:07:15 -0700 Subject: [PATCH] Type annotation: Scheduler.get_pareto_optimal_parameters never returns None (#2740) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/2740 Fixing a type annotation. Reviewed By: Balandat Differential Revision: D62192596 fbshipit-source-id: 4e5ce9ac4d9f54e98a69b60f5beba7afa3a50b09 --- ax/service/scheduler.py | 2 +- ax/service/utils/best_point_mixin.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/ax/service/scheduler.py b/ax/service/scheduler.py index 1b9066e7456..447ea39e437 100644 --- a/ax/service/scheduler.py +++ b/ax/service/scheduler.py @@ -469,7 +469,7 @@ def get_pareto_optimal_parameters( optimization_config: Optional[OptimizationConfig] = None, trial_indices: Optional[Iterable[int]] = None, use_model_predictions: bool = True, - ) -> Optional[dict[int, tuple[TParameterization, TModelPredictArm]]]: + ) -> dict[int, tuple[TParameterization, TModelPredictArm]]: return self._get_pareto_optimal_parameters( experiment=self.experiment, generation_strategy=self.standard_generation_strategy, diff --git a/ax/service/utils/best_point_mixin.py b/ax/service/utils/best_point_mixin.py index 80da762e578..75bfdf4eb34 100644 --- a/ax/service/utils/best_point_mixin.py +++ b/ax/service/utils/best_point_mixin.py @@ -134,7 +134,7 @@ def get_pareto_optimal_parameters( optimization_config: Optional[OptimizationConfig] = None, trial_indices: Optional[Iterable[int]] = None, use_model_predictions: bool = True, - ) -> Optional[dict[int, tuple[TParameterization, TModelPredictArm]]]: + ) -> dict[int, tuple[TParameterization, TModelPredictArm]]: """Identifies the best parameterizations tried in the experiment so far, using model predictions if ``use_model_predictions`` is true and using observed values from the experiment otherwise. By default, uses model @@ -158,12 +158,13 @@ def get_pareto_optimal_parameters( observed values. Returns: - ``None`` if it was not possible to extract the Pareto frontier, - otherwise a mapping from trial index to the tuple of: + A mapping from trial index to the tuple of: - the parameterization of the arm in that trial, - two-item tuple of metric means dictionary and covariance matrix (model-predicted if ``use_model_predictions=True`` and observed otherwise). + Raises a `NotImplementedError` if extracting the Pareto frontier is + not possible. Note that the returned dict may be empty. """ pass