lint with ufmt

Author: CompRhys

botorch/acquisition/input_constructors.py

@@ -1780,7 +1780,7 @@ def optimize_objective(
         bounds=free_feature_bounds,
         q=q,
         num_restarts=optimizer_options.get("num_restarts", 60),
-        raw_samples=optimizer_options.get("raw_samples", 1024),  # NOTE potential behaviour change
+        raw_samples=optimizer_options.get("raw_samples", 1024),
         options={
             "batch_limit": optimizer_options.get("batch_limit", 8),
             "maxiter": optimizer_options.get("maxiter", 200),

botorch/optim/optimize.py

@@ -113,8 +113,8 @@ def __post_init__(self) -> None:
                 # Use object.__setattr__ to bypass immutability and set a value
                 object.__setattr__(
                     self,
-                    'batch_initial_conditions',
-                    self.batch_initial_conditions.unsqueeze(0)
+                    "batch_initial_conditions",
+                    self.batch_initial_conditions.unsqueeze(0),
                 )
 
             if batch_initial_conditions_shape[-1] != d:
@@ -139,7 +139,7 @@ def __post_init__(self) -> None:
                     DeprecationWarning,
                 )
                 # Use object.__setattr__ to bypass immutability and set a value
-                object.__setattr__(self, 'raw_samples', None)
+                object.__setattr__(self, "raw_samples", None)
 
         elif self.ic_generator is None:
             if self.nonlinear_inequality_constraints is not None:
@@ -279,15 +279,20 @@ def _optimize_acqf_sequential_q(
     return candidates, torch.stack(acq_value_list)
 
 
-def _combine_initial_conditions(provided_initial_conditions, generated_initial_conditions, num_restarts):
-    if provided_initial_conditions is not None and generated_initial_conditions is not None:
+def _combine_initial_conditions(
+    provided_initial_conditions, generated_initial_conditions, num_restarts
+):
+    if (
+        provided_initial_conditions is not None
+        and generated_initial_conditions is not None
+    ):
         if (  # Repeat the provided initial conditions to match the number of restarts
             provided_initial_conditions.shape[0] == 1
             and num_restarts is not None
             and num_restarts > 1
         ):
             provided_initial_conditions = provided_initial_conditions.repeat(
-                num_restarts, *([1] * (provided_initial_conditions.dim()-1))
+                num_restarts, *([1] * (provided_initial_conditions.dim() - 1))
             )
         return torch.cat(
             [provided_initial_conditions, generated_initial_conditions], dim=-2
@@ -297,7 +302,9 @@ def _combine_initial_conditions(provided_initial_conditions, generated_initial_c
     elif generated_initial_conditions is not None:
         return generated_initial_conditions
     else:
-        raise ValueError("Either `batch_initial_conditions` or `raw_samples` must be set.")
+        raise ValueError(
+            "Either `batch_initial_conditions` or `raw_samples` must be set."
+        )
 
 
 def _optimize_acqf_batch(opt_inputs: OptimizeAcqfInputs) -> tuple[Tensor, Tensor]:
@@ -328,7 +335,7 @@ def _optimize_acqf_batch(opt_inputs: OptimizeAcqfInputs) -> tuple[Tensor, Tensor
     batch_initial_conditions = _combine_initial_conditions(
         provided_initial_conditions=opt_inputs.batch_initial_conditions,
         generated_initial_conditions=generated_initial_conditions,
-        num_restarts=opt_inputs.num_restarts
+        num_restarts=opt_inputs.num_restarts,
     )
 
     batch_limit: int = options.get(
@@ -429,7 +436,7 @@ def _optimize_batch_candidates() -> tuple[Tensor, Tensor, list[Warning]]:
             batch_initial_conditions = _combine_initial_conditions(
                 provided_initial_conditions=opt_inputs.batch_initial_conditions,
                 generated_initial_conditions=generated_initial_conditions,
-                num_restarts=opt_inputs.num_restarts
+                num_restarts=opt_inputs.num_restarts,
             )
 
             batch_candidates, batch_acq_values, ws = _optimize_batch_candidates()
@@ -1287,7 +1294,9 @@ def _gen_first_candidate_starting_points_local_search(
     generated_X0 = None
 
     if batch_initial_conditions is not None:
-        provided_X0 = _filter_invalid(X=batch_initial_conditions.squeeze(1), X_avoid=X_avoid)
+        provided_X0 = _filter_invalid(
+            X=batch_initial_conditions.squeeze(1), X_avoid=X_avoid
+        )
         provided_X0 = _filter_infeasible(
             X=provided_X0, inequality_constraints=inequality_constraints
         ).unsqueeze(1)
@@ -1312,7 +1321,9 @@ def _gen_first_candidate_starting_points_local_search(
     elif generated_X0 is not None:
         X0 = generated_X0
     else:
-        raise ValueError("Either `batch_initial_conditions` or `raw_samples` must be set.")
+        raise ValueError(
+            "Either `batch_initial_conditions` or `raw_samples` must be set."
+        )
 
     return X0
 

test/optim/test_optimize.py

@@ -546,7 +546,8 @@ def test_optimize_acqf_batch_limit(self) -> None:
 
         options = {"batch_limit": batch_limit}
         initial_conditions = [
-            torch.ones(shape) for shape in [(1, 2, dim), (3, 1, dim), (3, 6, dim), (1, dim)]
+            torch.ones(shape)
+            for shape in [(1, 2, dim), (3, 1, dim), (3, 6, dim), (1, dim)]
         ] + [None]
 
         for gen_candidates, ics in zip(
@@ -572,7 +573,9 @@ def test_optimize_acqf_batch_limit(self) -> None:
                 self.assertEqual(acq_value_list.shape, expected_shape)
 
         with self.subTest(gen_candidates=gen_candidates):
-            with self.assertWarnsRegex(DeprecationWarning, "Defaulting to old behavior"):
+            with self.assertWarnsRegex(
+                DeprecationWarning, "Defaulting to old behavior"
+            ):
                 ics = torch.ones((2, 1, dim))
                 _, acq_value_list = optimize_acqf(
                     acq_function=SinOneOverXAcqusitionFunction(),
@@ -602,7 +605,10 @@ def test_optimize_acqf_runs_given_batch_initial_conditions(self):
         ]
         q = 1
 
-        ic_shapes = [(1, 2, dim), (1, dim)]  # NOTE removed test on (2, 1, dim) as error tested elsewhere
+        ic_shapes = [
+            (1, 2, dim),
+            (1, dim),
+        ]  # NOTE removed test on (2, 1, dim) as error tested elsewhere
 
         torch.manual_seed(0)
         for shape in ic_shapes: