Clean up nits across acquisition module (meta-pytorch#3210)

Summary:
Pull Request resolved: meta-pytorch#3210

Fix code quality issues across multiple files in `botorch/acquisition/`:

**penalized.py:**
1. Rename misleading `pdf` variable to `penalty` in `GaussianPenalty.forward`.
2. Replace `__call__` with `forward` in `L0Approximation`, `L0PenaltyApprox`, and `L0PenaltyApproxObjective`.
3. Fix copy-paste docstring in `PenalizedAcquisitionFunction.__init__`.
4. Remove duplicated line in `PenalizedMCObjective` docstring example.
5. Replace bare `assert` with `ValueError` in `PenalizedMCObjective.forward`.

**fixed_feature.py:**
6. Fix return type annotation: `get_device_of_sequence` returns `torch.device`, not `torch.dtype`.
7. Fix `_is_cuda` to work on multi-GPU: use `value.device.type == "cuda"` instead of comparing against `torch.device("cuda")`.

**objective.py:**
8. Fix typo "Trasform" → "Transform" in comment.
9. Use `isinstance(eta, Tensor)` instead of `type(eta) is not Tensor`.
10. Replace bare `assert` with `ValueError` in `LearnedObjective.__init__`.

**proximal.py:**
11. Fix typos "decrese" → "decrease" and "vaidation" → "validation".

**thompson_sampling.py:**
12. Fix typo "inherting" → "inheriting".

**active_learning.py:**
13. Remove duplicate phrase "points that have" in `X_pending` docstring.
14. Fix return shape docstring: `batch_size` not `batch_size x q`.

**multioutput_acquisition.py:**
15. Fix `MultiOutputPosteriorMean` docstring: lists correct params (`model`, `weights`) instead of `acqfs`.
16. Fix `MultiOutputAcquisitionFunctionWrapper` docstring: says "wrapper" not "base class".

Reviewed By: dme65

Differential Revision: D94963222

fbshipit-source-id: 1a09876ab0aebb00e908cd5d5373dcafdc34723b

Author: saitcakmak

botorch/acquisition/active_learning.py

@@ -78,7 +78,7 @@ def __init__(
                 a PosteriorTransform that transforms the multi-output posterior into a
                 single-output posterior is required.
             X_pending: A ``n' x d``-dim Tensor of ``n'`` design points that have
-                points that have been submitted for function evaluation but
+                been submitted for function evaluation but
                 have not yet been evaluated.
         """
         super().__init__(model=model)
@@ -170,7 +170,7 @@ def forward(self, X: Tensor) -> Tensor:
             X: A ``batch_size x q x d``-dim Tensor. q should be a multiple of 2.
 
         Returns:
-            Tensor of shape ``batch_size x q`` representing the posterior variance
+            Tensor of shape ``batch_size`` representing the posterior variance
             of link function at X that active learning hopes to maximize
         """
         if X.shape[-2] == 0 or X.shape[-2] % 2 != 0:

botorch/acquisition/fixed_feature.py

@@ -35,15 +35,15 @@ def _is_single(value: Tensor | float) -> bool:
     return torch.float32 if all_single_precision else torch.float64
 
 
-def get_device_of_sequence(values: Sequence[Tensor | float]) -> torch.dtype:
+def get_device_of_sequence(values: Sequence[Tensor | float]) -> torch.device:
     """
     CPU if everything is on the CPU; Cuda otherwise.
 
     Numbers (non-tensors) are considered to be on the CPU.
     """
 
     def _is_cuda(value: Tensor | float) -> bool:
-        return hasattr(value, "device") and value.device == torch.device("cuda")
+        return isinstance(value, Tensor) and value.device.type == "cuda"
 
     any_cuda = any(_is_cuda(value) for value in values)
     return torch.device("cuda") if any_cuda else torch.device("cpu")

botorch/acquisition/multioutput_acquisition.py

@@ -56,13 +56,13 @@ def set_X_pending(self, X_pending: Tensor | None) -> None:
 
 class MultiOutputPosteriorMean(MultiOutputAcquisitionFunction):
     def __init__(self, model: Model, weights: Tensor | None = None) -> None:
-        r"""Constructor for the MultiPosteriorMean.
+        r"""Constructor for the MultiOutputPosteriorMean.
 
         Maximization of all outputs is assumed by default. Minimizing outputs can
         be achieved by setting the corresponding weights to negative.
 
         Args:
-            acqfs: A list of ``m`` acquisition functions.
+            model: A fitted model.
             weights: A one-dimensional tensor with ``m`` elements representing the
                 weights on the outputs.
         """
@@ -102,7 +102,7 @@ class MultiOutputAcquisitionFunctionWrapper(MultiOutputAcquisitionFunction):
     r"""Multi-output wrapper around single-output acquisition functions."""
 
     def __init__(self, acqfs: list[AcquisitionFunction]) -> None:
-        r"""Constructor for the AcquisitionFunction base class.
+        r"""Constructor for the MultiOutputAcquisitionFunctionWrapper.
 
         Args:
             acqfs: A list of ``m`` acquisition functions.

botorch/acquisition/objective.py

@@ -219,7 +219,7 @@ def forward(
         # ``n_w`` sized diagonal. The ``m`` outcomes are not interleaved.
         for i in range(q * m):
             weights[i, self.n_w * i : self.n_w * (i + 1)] = self.weights[:, i // q]
-        # Trasform the mean.
+        # Transform the mean.
         new_loc = (
             (weights @ org_mvn.loc.unsqueeze(-1))
             .view(*batch_shape, m, q)
@@ -452,7 +452,7 @@ def __init__(
         """
         super().__init__(objective=objective)
         self.constraints = constraints
-        if type(eta) is not Tensor:
+        if not isinstance(eta, Tensor):
             eta = torch.full((len(constraints),), eta)
         self.register_buffer("eta", eta)
         self.register_buffer("infeasible_cost", torch.as_tensor(infeasible_cost))
@@ -528,7 +528,8 @@ def __init__(
         super().__init__()
         self.pref_model = pref_model
         if isinstance(pref_model, DeterministicModel):
-            assert sample_shape is None
+            if sample_shape is not None:
+                raise ValueError("sample_shape must be None for DeterministicModel.")
             self.sampler = None
         else:
             if sample_shape is None:

botorch/acquisition/penalized.py

@@ -104,8 +104,8 @@ def forward(self, X: Tensor) -> Tensor:
             A tensor of size "batch_shape" representing the acqfn for each q-batch.
         """
         sq_diff = torch.linalg.norm((X - self.init_point), ord=2, dim=-1) ** 2
-        pdf = torch.exp(sq_diff / 2 / self.sigma**2)
-        regularization_term = pdf.max(dim=-1).values
+        penalty = torch.exp(sq_diff / 2 / self.sigma**2)
+        regularization_term = penalty.max(dim=-1).values
         return regularization_term
 
 
@@ -174,7 +174,7 @@ def __init__(self, target_point: Tensor, a: float = 1.0, **tkwargs: Any) -> None
         # hyperparameter to control the differentiable relaxation in L0 norm function.
         self.register_buffer("a", torch.tensor(a, **tkwargs))
 
-    def __call__(self, X: Tensor) -> Tensor:
+    def forward(self, X: Tensor) -> Tensor:
         return nnz_approx(X=X, target_point=self.target_point, a=self.a)
 
 
@@ -191,14 +191,14 @@ def __init__(self, target_point: Tensor, a: float = 1.0, **tkwargs: Any) -> None
         """
         super().__init__(target_point=target_point, a=a, **tkwargs)
 
-    def __call__(self, X: Tensor) -> Tensor:
+    def forward(self, X: Tensor) -> Tensor:
         r"""
         Args:
             X: A "batch_shape x q x dim" representing the points to be evaluated.
         Returns:
             A tensor of size "batch_shape" representing the acqfn for each q-batch.
         """
-        return super().__call__(X=X).squeeze(dim=-1).min(dim=-1).values
+        return super().forward(X=X).squeeze(dim=-1).min(dim=-1).values
 
 
 class PenalizedAcquisitionFunction(AcquisitionFunction):
@@ -216,7 +216,7 @@ def __init__(
         penalty_func: torch.nn.Module,
         regularization_parameter: float,
     ) -> None:
-        r"""Initializing Group-Lasso regularization.
+        r"""Initializing penalized acquisition function.
 
         Args:
             raw_acqf: The raw acquisition function that is going to be regularized.
@@ -317,7 +317,6 @@ class PenalizedMCObjective(GenericMCObjective):
                 objective, l1_penalty_objective, regularization_parameter
             )
         >>> samples = sampler(posterior)
-                objective, l1_penalty_objective, regularization_parameter
     """
 
     def __init__(
@@ -371,7 +370,12 @@ def forward(self, samples: Tensor, X: Tensor | None = None) -> Tensor:
         # tensor; obj returned from GenericMCObjective is a ``q``-dim tensor and
         # penalty_obj is a ``1 x q``-dim tensor.
         if obj.ndim == 1:
-            assert penalty_obj.shape == torch.Size([1, samples.shape[-2]])
+            if penalty_obj.shape != torch.Size([1, samples.shape[-2]]):
+                raise ValueError(  # pragma: no cover
+                    f"Expected penalty_obj of shape "
+                    f"{torch.Size([1, samples.shape[-2]])}, but got "
+                    f"{penalty_obj.shape}."
+                )
             penalty_obj = penalty_obj.squeeze(dim=0)
         return obj - self.regularization_parameter * penalty_obj
 
@@ -391,12 +395,12 @@ def __init__(self, target_point: Tensor, a: float = 1.0, **tkwargs: Any) -> None
         """
         super().__init__(target_point=target_point, a=a, **tkwargs)
 
-    def __call__(self, X: Tensor) -> Tensor:
+    def forward(self, X: Tensor) -> Tensor:
         r"""
         Args:
             X: A "batch_shape x q x dim" representing the points to be evaluated.
         Returns:
             A "1 x batch_shape x q" tensor representing the penalty for each point.
             The first dimension corresponds to the dimension of MC samples.
         """
-        return super().__call__(X=X).squeeze(dim=-1).unsqueeze(dim=0)
+        return super().forward(X=X).squeeze(dim=-1).unsqueeze(dim=0)

botorch/acquisition/proximal.py

@@ -34,7 +34,7 @@ class ProximalAcquisitionFunction(AcquisitionFunction):
     proximal weighting.
 
     Small values of ``proximal_weights`` corresponds to strong biasing towards recently
-    observed points, which smoothes optimization with a small potential decrese in
+    observed points, which smoothes optimization with a small potential decrease in
     convergence rate.
 
 
@@ -156,7 +156,7 @@ def forward(self, X: Tensor) -> Tensor:
 def _validate_model(model: Model, proximal_weights: Tensor) -> None:
     r"""Validate model
 
-    Perform vaidation checks on model used in base acquisition function to make sure
+    Perform validation checks on model used in base acquisition function to make sure
     it is compatible with proximal weighting.
 
     Args:

botorch/acquisition/thompson_sampling.py

@@ -60,7 +60,7 @@ def __init__(
         self.ensemble_indices: Tensor | None = None
 
         # NOTE: This conditional block is copied from MCAcquisitionFunction, we should
-        # consider inherting from it and e.g. getting the X_pending logic as well.
+        # consider inheriting from it and e.g. getting the X_pending logic as well.
         if objective is None and model.num_outputs != 1:
             if posterior_transform is None:
                 raise UnsupportedError(

test/acquisition/test_objective.py

@@ -570,7 +570,7 @@ def test_learned_preference_objective(self) -> None:
             self.assertAllClose(avg_obj_val, flipped_avg_obj_val)
 
         # cannot use a deterministic model together with a sampler
-        with self.subTest("deterministic model"), self.assertRaises(AssertionError):
+        with self.subTest("deterministic model"), self.assertRaises(ValueError):
             LearnedObjective(
                 pref_model=mean_pref_model,
                 sample_shape=torch.Size([num_samples]),