Update pyfmt component on FBS:master (facebook#4572)

Summary:
X-link: meta-pytorch/tritonbench#661

Pull Request resolved: facebook#4572

X-link: facebookexternal/aepsych_prerelease#42

X-link: facebook/dotslash#89

X-link: meta-pytorch/torchx#1165

X-link: facebookresearch/aepsych#824

X-link: meta-pytorch/botorch#3088

Differential Revision: D87671961

Author: bowiechen

.pre-commit-config.yaml

@@ -11,13 +11,13 @@ repos:
             - PyYAML
 
 -   repo: https://github.com/omnilib/ufmt
-    rev: v2.8.0
+    rev: v2.9.0
     hooks:
     -   id: ufmt
         additional_dependencies:
-        - black==24.4.2
-        - usort==1.0.8.post1
-        - ruff-api==0.1.0
+        - black==25.11.0
+        - usort==1.1.0
+        - ruff-api==0.2.0
         - stdlibs==2024.1.28
         args: [format]
 

ax/adapter/tests/test_base_adapter.py

@@ -618,10 +618,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(),
@@ -659,9 +662,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(),
@@ -1131,9 +1137,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

@@ -138,16 +138,19 @@ 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
-        with 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.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(model=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)

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

@@ -205,9 +205,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
@@ -318,9 +321,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
@@ -413,11 +417,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
             ):
@@ -814,14 +821,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

@@ -642,15 +642,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

@@ -732,9 +732,10 @@ 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, (
+                "Unexpected ordering of outcomes: "
+                f"{ordered_outcomes} != {self.outcomes}"
+            )
         return datasets, candidate_metadata, search_space_digest
 
     def _fit(

ax/adapter/transforms/stratified_standardize_y.py

@@ -124,8 +124,7 @@ def __init__(
         experiment_data = none_throws(experiment_data)
         if len(experiment_data.observation_data.index.names) > 2:
             raise NotImplementedError(
-                "StratifiedStandardizeY does not support experiment data with "
-                "map keys."
+                "StratifiedStandardizeY does not support experiment data with map keys."
             )
         strata = (
             experiment_data.arm_data[self.p_name]

ax/adapter/transforms/task_encode.py

@@ -40,9 +40,9 @@ def __init__(
         adapter: Optional["adapter_module.base.Adapter"] = None,
         config: TConfig | None = None,
     ) -> None:
-        assert (
-            search_space is not None
-        ), "TaskChoiceToIntTaskChoice requires search space"
+        assert search_space is not None, (
+            "TaskChoiceToIntTaskChoice requires search space"
+        )
         super().__init__(
             search_space=search_space,
             experiment_data=experiment_data,

ax/adapter/transforms/tests/test_bilog_y.py

@@ -13,7 +13,6 @@
 from itertools import product
 
 import numpy as np
-
 from ax.adapter.base import Adapter, DataLoaderConfig
 from ax.adapter.data_utils import extract_experiment_data
 from ax.adapter.transforms.bilog_y import bilog_transform, BilogY, inv_bilog_transform