apply Black 25.11.0 style in fbcode (77/92)

Summary:
Formats the covered files with pyfmt.

paintitblack

Reviewed By: itamaro

Differential Revision: D90476253

fbshipit-source-id: 1d5d238886f1ef4dd75646fede3890699740aa3c

Author: bowiechen

ax/adapter/tests/test_base_adapter.py

@@ -619,10 +619,13 @@ def test_set_status_quo_with_multiple_observations(self) -> None:
                 # Fetch constraint metric an additional time. This will lead to two
                 # separate observations for the status quo arm.
                 exp.fetch_data(metrics=[exp.metrics["branin_map_constraint"]])
-            with self.assertNoLogs(logger=logger, level="WARN"), mock.patch(
-                "ax.adapter.base._combine_multiple_status_quo_observations",
-                wraps=_combine_multiple_status_quo_observations,
-            ) as mock_combine:
+            with (
+                self.assertNoLogs(logger=logger, level="WARN"),
+                mock.patch(
+                    "ax.adapter.base._combine_multiple_status_quo_observations",
+                    wraps=_combine_multiple_status_quo_observations,
+                ) as mock_combine,
+            ):
                 adapter = Adapter(
                     experiment=exp,
                     generator=Generator(),
@@ -660,9 +663,12 @@ def test_set_status_quo_with_multiple_observations(self) -> None:
             )
 
         # Case 2: Experiment has an optimization config with no map metrics
-        with mock.patch(
-            "ax.adapter.base.has_map_metrics", return_value=False
-        ) as mock_extract, self.assertLogs(logger=logger, level="WARN") as mock_logs:
+        with (
+            mock.patch(
+                "ax.adapter.base.has_map_metrics", return_value=False
+            ) as mock_extract,
+            self.assertLogs(logger=logger, level="WARN") as mock_logs,
+        ):
             adapter = Adapter(
                 experiment=exp,
                 generator=Generator(),
@@ -1109,9 +1115,10 @@ def mock_predict(
         self.assertTrue(np.allclose(f["m1"], np.ones(3) * 2.0))
 
         # Test for error if an observation is dropped.
-        with mock.patch.object(
-            adapter, "_predict", side_effect=mock_predict
-        ), self.assertRaisesRegex(ModelError, "Predictions resulted in fewer"):
+        with (
+            mock.patch.object(adapter, "_predict", side_effect=mock_predict),
+            self.assertRaisesRegex(ModelError, "Predictions resulted in fewer"),
+        ):
             adapter.predict(
                 observation_features=[
                     ObservationFeatures(parameters={"x": 3.0, "y": 4.0}),

ax/adapter/tests/test_cross_validation.py

@@ -141,12 +141,15 @@ def test_cross_validate_base(self) -> None:
         )
 
         # Test LOO - use naive CV path by mocking efficient LOO
-        with mock.patch(
-            "ax.adapter.cross_validation._efficient_loo_cross_validate",
-            side_effect=ValueError("Force fallback to naive CV"),
-        ), mock.patch.object(
-            self.adapter, "cross_validate", wraps=self.adapter.cross_validate
-        ) as mock_cv:
+        with (
+            mock.patch(
+                "ax.adapter.cross_validation._efficient_loo_cross_validate",
+                side_effect=ValueError("Force fallback to naive CV"),
+            ),
+            mock.patch.object(
+                self.adapter, "cross_validate", wraps=self.adapter.cross_validate
+            ) as mock_cv,
+        ):
             result = cross_validate(adapter=self.adapter, folds=-1)
         self.assertEqual(len(result), 4)
         z = mock_cv.mock_calls
@@ -163,19 +166,23 @@ def test_cross_validate_base(self) -> None:
             np.array_equal(sorted(all_test), np.array([2.0, 2.0, 3.0, 4.0]))
         )
         # Test LOO in transformed space - use naive path by mocking efficient LOO
-        with mock.patch(
-            "ax.adapter.cross_validation._efficient_loo_cross_validate",
-            side_effect=ValueError("Force fallback to naive CV"),
-        ), mock.patch.object(
-            self.adapter,
-            "_transform_inputs_for_cv",
-            wraps=self.adapter._transform_inputs_for_cv,
-        ) as mock_transform_cv, mock.patch.object(
-            self.adapter,
-            "_cross_validate",
-            side_effect=lambda **kwargs: [self.observation_data]
-            * len(kwargs["cv_test_points"]),
-        ) as mock_cv:
+        with (
+            mock.patch(
+                "ax.adapter.cross_validation._efficient_loo_cross_validate",
+                side_effect=ValueError("Force fallback to naive CV"),
+            ),
+            mock.patch.object(
+                self.adapter,
+                "_transform_inputs_for_cv",
+                wraps=self.adapter._transform_inputs_for_cv,
+            ) as mock_transform_cv,
+            mock.patch.object(
+                self.adapter,
+                "_cross_validate",
+                side_effect=lambda **kwargs: [self.observation_data]
+                * len(kwargs["cv_test_points"]),
+            ) as mock_cv,
+        ):
             result = cross_validate(adapter=self.adapter, folds=-1, untransform=False)
         result_predicted_obs_data = [cv_result.predicted for cv_result in result]
         self.assertEqual(result_predicted_obs_data, [self.observation_data] * 4)
@@ -246,12 +253,15 @@ def test_selector(obs: Observation) -> bool:
 
         # test observation noise - use naive path by disabling efficient LOO
         for untransform in (True, False):
-            with mock.patch(
-                "ax.adapter.cross_validation._efficient_loo_cross_validate",
-                side_effect=ValueError("Force fallback to naive CV"),
-            ), mock.patch.object(
-                self.adapter, "_cross_validate", wraps=self.adapter._cross_validate
-            ) as mock_cv:
+            with (
+                mock.patch(
+                    "ax.adapter.cross_validation._efficient_loo_cross_validate",
+                    side_effect=ValueError("Force fallback to naive CV"),
+                ),
+                mock.patch.object(
+                    self.adapter, "_cross_validate", wraps=self.adapter._cross_validate
+                ) as mock_cv,
+            ):
                 result = cross_validate(
                     adapter=self.adapter,
                     folds=-1,
@@ -500,9 +510,12 @@ def test_has_good_opt_config_model_fit(self) -> None:
     def test_efficient_loo_cv_is_attempted(self) -> None:
         """Test that efficient LOO CV is attempted only when all conditions are met."""
         # Setup adapter with a BoTorchGenerator
-        with mock.patch(
-            "botorch.cross_validation.efficient_loo_cv"
-        ) as mock_efficient_loo, mock.patch("botorch.cross_validation.ensemble_loo_cv"):
+        with (
+            mock.patch(
+                "botorch.cross_validation.efficient_loo_cv"
+            ) as mock_efficient_loo,
+            mock.patch("botorch.cross_validation.ensemble_loo_cv"),
+        ):
             # Create mock LOO results
             # Create a mock posterior
             mock_mean = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
@@ -570,11 +583,13 @@ def _fold_gen(td: ExperimentData) -> Iterable[CVData]:
 
         # For adapter with aux experiments, directly verify the condition check
         # rather than running through the full cross_validate path
-        with self.subTest(condition="has auxiliary experiments"), mock.patch(
-            "ax.adapter.cross_validation._efficient_loo_cross_validate"
-        ) as mock_efficient, mock.patch(
-            "ax.adapter.cross_validation._fold_cross_validate"
-        ) as mock_fold:
+        with (
+            self.subTest(condition="has auxiliary experiments"),
+            mock.patch(
+                "ax.adapter.cross_validation._efficient_loo_cross_validate"
+            ) as mock_efficient,
+            mock.patch("ax.adapter.cross_validation._fold_cross_validate") as mock_fold,
+        ):
             mock_fold.return_value = []
             cross_validate(adapter=adapter_with_aux)
             self.assertFalse(
@@ -584,9 +599,12 @@ def _fold_gen(td: ExperimentData) -> Iterable[CVData]:
 
         for kwargs, adapter_override, desc in conditions_preventing_efficient_loo:
             adapter = adapter_override or self.adapter
-            with self.subTest(condition=desc), mock.patch(
-                "ax.adapter.cross_validation._efficient_loo_cross_validate"
-            ) as mock_efficient:
+            with (
+                self.subTest(condition=desc),
+                mock.patch(
+                    "ax.adapter.cross_validation._efficient_loo_cross_validate"
+                ) as mock_efficient,
+            ):
                 # pyre-ignore[6]: kwargs is properly typed for cross_validate
                 cross_validate(adapter=adapter, **kwargs)
                 self.assertFalse(
@@ -596,13 +614,15 @@ def _fold_gen(td: ExperimentData) -> Iterable[CVData]:
 
         # Test logger when efficient LOO fails even though all conditions were met
         with self.subTest(condition="efficient LOO fails with exception"):
-            with mock.patch(
-                "ax.adapter.cross_validation._efficient_loo_cross_validate"
-            ) as mock_efficient, mock.patch(
-                "ax.adapter.cross_validation._fold_cross_validate"
-            ) as mock_fold, mock.patch(
-                "ax.adapter.cross_validation.logger"
-            ) as mock_logger:
+            with (
+                mock.patch(
+                    "ax.adapter.cross_validation._efficient_loo_cross_validate"
+                ) as mock_efficient,
+                mock.patch(
+                    "ax.adapter.cross_validation._fold_cross_validate"
+                ) as mock_fold,
+                mock.patch("ax.adapter.cross_validation.logger") as mock_logger,
+            ):
                 # Force efficient LOO to fail
                 mock_efficient.side_effect = ValueError("Test failure reason")
                 mock_fold.return_value = []
@@ -701,13 +721,16 @@ def _test_efficient_loo_cv_matches_naive(
             )
 
         # Run naive CV (by forcing fallback)
-        with mock.patch(
-            "ax.adapter.cross_validation._efficient_loo_cross_validate",
-            side_effect=ValueError("Force fallback to naive CV"),
-        ), mock.patch(
-            "ax.adapter.cross_validation._fold_cross_validate",
-            wraps=_fold_cross_validate,
-        ) as mock_naive_cv:
+        with (
+            mock.patch(
+                "ax.adapter.cross_validation._efficient_loo_cross_validate",
+                side_effect=ValueError("Force fallback to naive CV"),
+            ),
+            mock.patch(
+                "ax.adapter.cross_validation._fold_cross_validate",
+                wraps=_fold_cross_validate,
+            ) as mock_naive_cv,
+        ):
             result_naive = cross_validate(
                 adapter=adapter,
                 folds=-1,
@@ -719,12 +742,15 @@ def _test_efficient_loo_cv_matches_naive(
             self.assertTrue(mock_naive_cv.called, "Naive CV not called")
 
         # Run efficient CV
-        with mock.patch(
-            "ax.adapter.cross_validation._efficient_loo_cross_validate",
-            wraps=_efficient_loo_cross_validate,
-        ) as mock_efficient, mock.patch(
-            "ax.adapter.cross_validation._fold_cross_validate",
-        ) as mock_naive:
+        with (
+            mock.patch(
+                "ax.adapter.cross_validation._efficient_loo_cross_validate",
+                wraps=_efficient_loo_cross_validate,
+            ) as mock_efficient,
+            mock.patch(
+                "ax.adapter.cross_validation._fold_cross_validate",
+            ) as mock_naive,
+        ):
             result_efficient = cross_validate(
                 adapter=adapter,
                 folds=-1,
@@ -765,14 +791,12 @@ def sort_key(cv_result: CVResult) -> tuple[float, ...]:
                 if untransform:
                     self.assertTrue(
                         np.all(obs_means > 5.0),
-                        f"untransform=True: expected original space, "
-                        f"got {obs_means}",
+                        f"untransform=True: expected original space, got {obs_means}",
                     )
                 else:
                     self.assertTrue(
                         np.all(np.abs(obs_means) < 3.0),
-                        f"untransform=False: expected standardized, "
-                        f"got {obs_means}",
+                        f"untransform=False: expected standardized, got {obs_means}",
                     )
 
         # Compare predictions

ax/adapter/tests/test_hierarchical_search_space.py

@@ -12,7 +12,6 @@
 
 from ax.adapter.cross_validation import cross_validate
 from ax.adapter.registry import Generators
-
 from ax.core.experiment import Experiment
 from ax.core.objective import Objective
 from ax.core.observation import ObservationFeatures
@@ -23,7 +22,6 @@
     ParameterType,
     RangeParameter,
 )
-
 from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
 from ax.metrics.noisy_function import GenericNoisyFunctionMetric

ax/adapter/tests/test_torch_adapter.py

@@ -214,9 +214,12 @@ def test_TorchAdapter(self, device: torch.device | None = None) -> None:
         pending_observations = {
             "y2": [ObservationFeatures(parameters={"x1": 1.0, "x2": 2.0, "x3": 3.0})]
         }
-        with ExitStack() as es, mock.patch.object(
-            generator, "gen", return_value=gen_return_value
-        ) as mock_gen:
+        with (
+            ExitStack() as es,
+            mock.patch.object(
+                generator, "gen", return_value=gen_return_value
+            ) as mock_gen,
+        ):
             es.enter_context(
                 mock.patch.object(
                     generator, "best_point", return_value=best_point_return_value
@@ -327,9 +330,10 @@ def test_evaluate_acquisition_function(self) -> None:
         obsf = ObservationFeatures(parameters={"x1": 1.0, "x2": 2.0})
 
         # Check for value error when optimization config is not set.
-        with mock.patch.object(
-            adapter, "_optimization_config", None
-        ), self.assertRaisesRegex(ValueError, "optimization_config"):
+        with (
+            mock.patch.object(adapter, "_optimization_config", None),
+            self.assertRaisesRegex(ValueError, "optimization_config"),
+        ):
             adapter.evaluate_acquisition_function(observation_features=[obsf])
 
         mock_acq_val = 5.0
@@ -422,11 +426,14 @@ def test_best_point(self) -> None:
         gen_return_value = TorchGenResults(
             points=torch.tensor([[1.0]]), weights=torch.tensor([1.0])
         )
-        with mock.patch(
-            f"{TorchGenerator.__module__}.TorchGenerator.best_point",
-            return_value=torch.tensor([best_point_value]),
-            autospec=True,
-        ), mock.patch.object(adapter, "predict", return_value=predict_return_value):
+        with (
+            mock.patch(
+                f"{TorchGenerator.__module__}.TorchGenerator.best_point",
+                return_value=torch.tensor([best_point_value]),
+                autospec=True,
+            ),
+            mock.patch.object(adapter, "predict", return_value=predict_return_value),
+        ):
             with mock.patch.object(
                 adapter.generator, "gen", return_value=gen_return_value
             ):
@@ -823,14 +830,17 @@ def test_gen_metadata_untransform(self) -> None:
                 weights=torch.tensor([1.0]),
                 gen_metadata={Keys.EXPECTED_ACQF_VAL: [1.0], **additional_metadata},
             )
-            with mock.patch.object(
-                adapter,
-                "_untransform_objective_thresholds",
-                wraps=adapter._untransform_objective_thresholds,
-            ) as mock_untransform, mock.patch.object(
-                generator,
-                "gen",
-                return_value=gen_return_value,
+            with (
+                mock.patch.object(
+                    adapter,
+                    "_untransform_objective_thresholds",
+                    wraps=adapter._untransform_objective_thresholds,
+                ) as mock_untransform,
+                mock.patch.object(
+                    generator,
+                    "gen",
+                    return_value=gen_return_value,
+                ),
             ):
                 adapter.gen(n=1)
             if additional_metadata.get("objective_thresholds", None) is None:

ax/adapter/tests/test_torch_moo_adapter.py

@@ -640,15 +640,18 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
             torch_device=torch.device("cuda" if cuda else "cpu"),
         )
         self.assertIn("Cast", adapter.transforms)
-        with patch.object(
-            adapter,
-            "_untransform_objective_thresholds",
-            wraps=adapter._untransform_objective_thresholds,
-        ) as mock_untransform, patch.object(
-            adapter.transforms["Cast"],
-            "untransform_observation_features",
-            wraps=adapter.transforms["Cast"].untransform_observation_features,
-        ) as wrapped_cast:
+        with (
+            patch.object(
+                adapter,
+                "_untransform_objective_thresholds",
+                wraps=adapter._untransform_objective_thresholds,
+            ) as mock_untransform,
+            patch.object(
+                adapter.transforms["Cast"],
+                "untransform_observation_features",
+                wraps=adapter.transforms["Cast"].untransform_observation_features,
+            ) as wrapped_cast,
+        ):
             obj_thresholds = adapter.infer_objective_thresholds(
                 search_space=exp.search_space,
                 optimization_config=exp.optimization_config,

ax/adapter/torch.py

@@ -780,9 +780,9 @@ def _get_fit_args(
         if update_outcomes_and_parameters:
             self.outcomes = ordered_outcomes
         else:
-            assert (
-                ordered_outcomes == self.outcomes
-            ), f"Unexpected ordering of outcomes: {ordered_outcomes} != {self.outcomes}"
+            assert ordered_outcomes == self.outcomes, (
+                f"Unexpected ordering of outcomes: {ordered_outcomes} != {self.outcomes}"
+            )
         return datasets, candidate_metadata, search_space_digest
 
     def _fit(