Replace checked_cast with pyre_extensions assert_is_instance (#3229)

Summary:
Pull Request resolved: #3229

Replace `checked_cast` with pyre_extensions `assert_is_instance`

Search for all use cases with "ax.utils.common.typeutils .* checked_cast" regex

Reviewed By: esantorella

Differential Revision: D67879879

fbshipit-source-id: c009242e71aab2093d96577615b1c7c408440fd0

Author: eonofrey

ax/analysis/healthcheck/constraints_feasibility.py

@@ -26,8 +26,7 @@
 from ax.modelbridge.base import ModelBridge
 from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.transforms.derelativize import Derelativize
-from ax.utils.common.typeutils import checked_cast
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 
 class ConstraintsFeasibilityAnalysis(HealthcheckAnalysis):
@@ -101,12 +100,8 @@ def compute(
             raise UserInputError(
                 "ConstraintsFeasibilityAnalysis requires a GenerationStrategy."
             )
-        generation_strategy = checked_cast(
-            GenerationStrategy,
-            generation_strategy,
-            exception=UserInputError(
-                "ConstraintsFeasibilityAnalysis requires a GenerationStrategy."
-            ),
+        generation_strategy = assert_is_instance(
+            generation_strategy, GenerationStrategy
         )
 
         if generation_strategy.model is None:
@@ -120,8 +115,8 @@ def compute(
                 "The current model is {model._model_key} and does not support "
                 "prediction."
             )
-        optimization_config = checked_cast(
-            OptimizationConfig, experiment.optimization_config
+        optimization_config = assert_is_instance(
+            experiment.optimization_config, OptimizationConfig
         )
         constraints_feasible, df = constraints_feasibility(
             optimization_config=optimization_config,

ax/analysis/healthcheck/search_space_analysis.py

@@ -25,7 +25,7 @@
 from ax.core.search_space import SearchSpace
 from ax.core.types import TParameterization
 from ax.exceptions.core import UserInputError
-from ax.utils.common.typeutils import checked_cast
+from pyre_extensions import assert_is_instance
 
 
 class SearchSpaceAnalysis(HealthcheckAnalysis):
@@ -141,7 +141,9 @@ def search_space_boundary_proportions(
             lower = parameter.lower
             upper = parameter.upper
         elif isinstance(parameter, ChoiceParameter) and parameter.is_ordered:
-            values = [checked_cast(Union[int, float], v) for v in parameter.values]
+            values = [
+                assert_is_instance(v, Union[int, float]) for v in parameter.values
+            ]
             lower = min(values)
             upper = max(values)
         else:
@@ -176,7 +178,8 @@ def search_space_boundary_proportions(
     for pc in search_space.parameter_constraints:
         weighted_sums = [
             sum(
-                float(checked_cast(Union[int, float], parametrization[param])) * weight
+                float(assert_is_instance(parametrization[param], Union[int, float]))
+                * weight
                 for param, weight in pc.constraint_dict.items()
             )
             for parametrization in parametrizations

ax/analysis/healthcheck/tests/test_constraints_feasibility.py

@@ -26,10 +26,9 @@
 from ax.modelbridge.model_spec import ModelSpec
 from ax.modelbridge.registry import Models
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import checked_cast
 from ax.utils.testing.core_stubs import get_branin_experiment_with_multi_objective
 from ax.utils.testing.mock import mock_botorch_optimize
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 
 class TestConstraintsFeasibilityAnalysis(TestCase):
@@ -74,7 +73,7 @@ def setUp(self) -> None:
         sobol = get_sobol(search_space=experiment.search_space)
         experiment.new_batch_trial(generator_run=sobol.gen(5))
 
-        batch_trial = checked_cast(BatchTrial, experiment.trials[0])
+        batch_trial = assert_is_instance(experiment.trials[0], BatchTrial)
 
         batch_trial.add_arm(experiment.status_quo)
         batch_trial.set_status_quo_with_weight(
@@ -107,8 +106,8 @@ def setUp(self) -> None:
     def test_constraints_feasibility(self) -> None:
         self.setUp()
         model = none_throws(self.generation_strategy.model)
-        optimization_config = checked_cast(
-            OptimizationConfig, self.experiment.optimization_config
+        optimization_config = assert_is_instance(
+            self.experiment.optimization_config, OptimizationConfig
         )
         constraints_feasible, df_arms = constraints_feasibility(
             optimization_config=optimization_config,
@@ -136,8 +135,8 @@ def test_constraints_feasibility(self) -> None:
         experiment.attach_data(data=Data(df=df))
         generation_strategy._fit_current_model(data=experiment.lookup_data())
         model = none_throws(generation_strategy.model)
-        optimization_config = checked_cast(
-            OptimizationConfig, experiment.optimization_config
+        optimization_config = assert_is_instance(
+            experiment.optimization_config, OptimizationConfig
         )
         constraints_feasible, df_arms = constraints_feasibility(
             optimization_config=optimization_config, model=model
@@ -146,8 +145,8 @@ def test_constraints_feasibility(self) -> None:
         experiment.optimization_config = OptimizationConfig(
             objective=Objective(metric=Metric(name="branin_a"), minimize=False),
         )
-        optimization_config = checked_cast(
-            OptimizationConfig, experiment.optimization_config
+        optimization_config = assert_is_instance(
+            experiment.optimization_config, OptimizationConfig
         )
         with self.assertRaises(UserInputError):
             constraints_feasibility(

ax/analysis/plotly/arm_effects/predicted_effects.py

@@ -26,8 +26,7 @@
 from ax.modelbridge.base import ModelBridge
 from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.transforms.derelativize import Derelativize
-from ax.utils.common.typeutils import checked_cast
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 
 class PredictedEffectsPlot(PlotlyAnalysis):
@@ -73,14 +72,14 @@ def compute(
     ) -> PlotlyAnalysisCard:
         if experiment is None:
             raise UserInputError("PredictedEffectsPlot requires an Experiment.")
-
-        generation_strategy = checked_cast(
-            GenerationStrategy,
-            generation_strategy,
-            exception=UserInputError(
+        try:
+            generation_strategy = assert_is_instance(
+                generation_strategy, GenerationStrategy
+            )
+        except TypeError as e:
+            raise UserInputError(
                 "PredictedEffectsPlot requires a GenerationStrategy."
-            ),
-        )
+            ) from e
 
         try:
             trial_indices = [

ax/analysis/plotly/tests/test_cross_validation.py

@@ -11,9 +11,8 @@
 from ax.exceptions.core import UserInputError
 from ax.service.ax_client import AxClient, ObjectiveProperties
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import checked_cast
 from ax.utils.testing.mock import mock_botorch_optimize
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 
 class TestCrossValidationPlot(TestCase):
@@ -71,7 +70,7 @@ def test_compute(self) -> None:
             # and therefore hasn't observed it
             if t.index == max(self.client.experiment.trials.keys()):
                 continue
-            arm_name = none_throws(checked_cast(Trial, t).arm).name
+            arm_name = none_throws(assert_is_instance(t, Trial).arm).name
             self.assertIn(
                 arm_name,
                 card.df["arm_name"].unique(),
@@ -93,7 +92,7 @@ def test_it_can_specify_trial_index_correctly(self) -> None:
             # and therefore hasn't observed it
             if t.index == max(self.client.experiment.trials.keys()):
                 continue
-            arm_name = none_throws(checked_cast(Trial, t).arm).name
+            arm_name = none_throws(assert_is_instance(t, Trial).arm).name
             self.assertIn(
                 arm_name,
                 card.df["arm_name"].unique(),

ax/analysis/plotly/tests/test_predicted_effects.py

@@ -19,7 +19,6 @@
 from ax.modelbridge.prediction_utils import predict_at_point
 from ax.modelbridge.registry import Models
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import checked_cast
 from ax.utils.testing.core_stubs import (
     get_branin_experiment,
     get_branin_metric,
@@ -28,7 +27,7 @@
 from ax.utils.testing.mock import mock_botorch_optimize
 from ax.utils.testing.modeling_stubs import get_sobol_MBM_MTGP_gs
 from botorch.utils.probability.utils import compute_log_prob_feas_from_bounds
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 
 class TestPredictedEffectsPlot(TestCase):
@@ -298,7 +297,7 @@ def test_it_works_for_non_batch_experiments(self) -> None:
         # THEN it has all arms represented in the dataframe
         for trial in experiment.trials.values():
             self.assertIn(
-                none_throws(checked_cast(Trial, trial).arm).name,
+                none_throws(assert_is_instance(trial, Trial).arm).name,
                 card.df["arm_name"].unique(),
             )
 

ax/benchmark/tests/test_benchmark_problem.py

@@ -26,7 +26,6 @@
 from ax.core.types import ComparisonOp
 from ax.exceptions.core import UserInputError
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import checked_cast
 from botorch.test_functions.base import ConstrainedBaseTestProblem
 from botorch.test_functions.multi_objective import BraninCurrin, ConstrainedBraninCurrin
 from botorch.test_functions.synthetic import (
@@ -225,14 +224,14 @@ def _test_constrained_from_botorch(
         )
 
         self.assertEqual(
-            checked_cast(BenchmarkMetric, metric).observe_noise_sd,
+            assert_is_instance(metric, BenchmarkMetric).observe_noise_sd,
             observe_noise_sd,
         )
 
         # TODO: Support observing noise variance only for some outputs
         for constraint in outcome_constraints:
             self.assertEqual(
-                checked_cast(BenchmarkMetric, constraint.metric).observe_noise_sd,
+                assert_is_instance(constraint.metric, BenchmarkMetric).observe_noise_sd,
                 observe_noise_sd,
             )
 

ax/benchmark/tests/test_benchmark_runner.py

@@ -32,7 +32,6 @@
 from ax.core.trial import Trial
 from ax.exceptions.core import UnsupportedError
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import checked_cast
 from ax.utils.testing.benchmark_stubs import (
     DummyTestFunction,
     get_jenatton_trials,
@@ -42,7 +41,7 @@
 from botorch.test_functions.synthetic import Ackley, ConstrainedHartmann, Hartmann
 from botorch.utils.transforms import normalize
 from pandas import DataFrame
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 
 class TestBenchmarkRunner(TestCase):
@@ -315,7 +314,7 @@ def test_heterogeneous_noise(self) -> None:
                 noise_std=noise_std,
             )
             self.assertDictEqual(
-                checked_cast(dict, runner.get_noise_stds()), noise_dict
+                assert_is_instance(runner.get_noise_stds(), dict), noise_dict
             )
 
             X = torch.rand(1, 6, dtype=torch.double)

ax/core/batch_trial.py

@@ -34,8 +34,7 @@
 from ax.utils.common.docutils import copy_doc
 from ax.utils.common.equality import datetime_equals, equality_typechecker
 from ax.utils.common.logger import _round_floats_for_logging, get_logger
-from ax.utils.common.typeutils import checked_cast
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 logger: Logger = get_logger(__name__)
 
@@ -490,7 +489,10 @@ def is_factorial(self) -> bool:
         return len(self.arms) == param_cardinality
 
     def run(self) -> BatchTrial:
-        return checked_cast(BatchTrial, super().run())
+        return assert_is_instance(
+            super().run(),
+            BatchTrial,
+        )
 
     def normalized_arm_weights(
         self, total: float = 1, trunc_digits: int | None = None

ax/core/data.py

@@ -27,8 +27,7 @@
     TClassDecoderRegistry,
     TDecoderRegistry,
 )
-from ax.utils.common.typeutils import checked_cast
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 TBaseData = TypeVar("TBaseData", bound="BaseData")
 DF_REPR_MAX_LENGTH = 1000
@@ -148,7 +147,7 @@ def _safecast_df(
             and coltype is not Any
         }
 
-        return checked_cast(pd.DataFrame, df.astype(dtype=dtype))
+        return assert_is_instance(df.astype(dtype=dtype), pd.DataFrame)
 
     @classmethod
     def required_columns(cls) -> set[str]:
@@ -194,7 +193,7 @@ def serialize_init_args(cls, obj: Any) -> dict[str, Any]:
         """Serialize the class-dependent properties needed to initialize this Data.
         Used for storage and to help construct new similar Data.
         """
-        data = checked_cast(cls, obj)
+        data = assert_is_instance(obj, cls)
         return serialize_init_args(
             obj=data, exclude_fields=["_skip_ordering_and_validation"]
         )

ax/core/experiment.py

@@ -55,8 +55,7 @@
 from ax.utils.common.logger import _round_floats_for_logging, get_logger
 from ax.utils.common.result import Err, Ok
 from ax.utils.common.timeutils import current_timestamp_in_millis
-from ax.utils.common.typeutils import checked_cast
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 logger: logging.Logger = get_logger(__name__)
 
@@ -1389,8 +1388,8 @@ def warm_start_from_old_experiment(
                 old_data = (
                     old_experiment.default_data_constructor(
                         df=new_df,
-                        map_key_infos=checked_cast(
-                            MapData, old_experiment.lookup_data()
+                        map_key_infos=assert_is_instance(
+                            old_experiment.lookup_data(), MapData
                         ).map_key_infos,
                     )
                     if old_experiment.default_data_type == DataType.MAP_DATA
@@ -1603,7 +1602,7 @@ def attach_trial(
             # data for this arm "complete" in the flattened search space.
             candidate_metadata = None
             if self.search_space.is_hierarchical:
-                hss = checked_cast(HierarchicalSearchSpace, self.search_space)
+                hss = assert_is_instance(self.search_space, HierarchicalSearchSpace)
                 candidate_metadata = hss.cast_observation_features(
                     observation_features=hss.flatten_observation_features(
                         observation_features=observation.ObservationFeatures(
@@ -1785,7 +1784,9 @@ def metric_config_summary_df(self) -> pd.DataFrame:
         if self.optimization_config is not None:
             opt_config = self.optimization_config
             if self.is_moo_problem:
-                multi_objective = checked_cast(MultiObjective, opt_config.objective)
+                multi_objective = assert_is_instance(
+                    opt_config.objective, MultiObjective
+                )
                 objectives = multi_objective.objectives
             else:
                 objectives = [opt_config.objective]

ax/core/map_data.py

@@ -28,7 +28,7 @@
     TClassDecoderRegistry,
     TDecoderRegistry,
 )
-from ax.utils.common.typeutils import checked_cast
+from pyre_extensions import assert_is_instance
 
 logger: Logger = get_logger(__name__)
 
@@ -327,7 +327,7 @@ def filter(
     @classmethod
     # pyre-fixme[2]: Parameter annotation cannot be `Any`.
     def serialize_init_args(cls, obj: Any) -> dict[str, Any]:
-        map_data = checked_cast(MapData, obj)
+        map_data = assert_is_instance(obj, MapData)
         properties = serialize_init_args(
             obj=map_data, exclude_fields=["_skip_ordering_and_validation"]
         )

ax/core/observation.py

@@ -28,8 +28,7 @@
 from ax.utils.common.base import Base
 from ax.utils.common.constants import Keys
 from ax.utils.common.logger import get_logger
-from ax.utils.common.typeutils import checked_cast
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 
 logger: Logger = get_logger(__name__)
 
@@ -429,7 +428,7 @@ def get_feature_cols(data: Data, is_map_data: bool = False) -> list[str]:
     # use observations_from_map_data, which is required
     # to properly handle MapData features (e.g. fidelity).
     if is_map_data:
-        data = checked_cast(MapData, data)
+        data = assert_is_instance(data, MapData)
         feature_cols = feature_cols.union(data.map_keys)
 
     for column in TIME_COLS: