Remove unused ** arguments (#2327)

Summary:
Pull Request resolved: #2327

In the process of auditing places where arguments (especially of the ** form) are ignored in BoTorch, I fixed the occurrences that seemed easy.

Reviewed By: saitcakmak

Differential Revision: D56828338

botorch/acquisition/cost_aware.py

@@ -13,7 +13,7 @@
 
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Optional, Union
+from typing import Callable, Optional, Union
 
 import torch
 from botorch import settings
@@ -38,7 +38,9 @@ class CostAwareUtility(Module, ABC):
     """
 
     @abstractmethod
-    def forward(self, X: Tensor, deltas: Tensor, **kwargs: Any) -> Tensor:
+    def forward(
+        self, X: Tensor, deltas: Tensor, sampler: Optional[MCSampler] = None
+    ) -> Tensor:
         r"""Evaluate the cost-aware utility on the candidates and improvements.
 
         Args:
@@ -47,6 +49,8 @@ def forward(self, X: Tensor, deltas: Tensor, **kwargs: Any) -> Tensor:
             deltas: A `num_fantasies x batch_shape`-dim Tensor of `num_fantasy`
                 samples from the marginal improvement in utility over the
                 current state at `X` for each t-batch.
+            sampler: A sampler used for sampling from the posterior of the cost
+                model. Some subclasses ignore this argument.
 
         Returns:
             A `num_fantasies x batch_shape`-dim Tensor of cost-transformed utilities.
@@ -67,7 +71,9 @@ def __init__(self, cost: Callable[[Tensor, Tensor], Tensor]) -> None:
         super().__init__()
         self._cost_callable: Callable[[Tensor, Tensor], Tensor] = cost
 
-    def forward(self, X: Tensor, deltas: Tensor, **kwargs: Any) -> Tensor:
+    def forward(
+        self, X: Tensor, deltas: Tensor, sampler: Optional[MCSampler] = None
+    ) -> Tensor:
         r"""Evaluate the cost function on the candidates and improvements.
 
         Args:
@@ -76,6 +82,7 @@ def forward(self, X: Tensor, deltas: Tensor, **kwargs: Any) -> Tensor:
             deltas: A `num_fantasies x batch_shape`-dim Tensor of `num_fantasy`
                 samples from the marginal improvement in utility over the
                 current state at `X` for each t-batch.
+            sampler: Ignored.
 
         Returns:
             A `num_fantasies x batch_shape`-dim Tensor of cost-weighted utilities.
@@ -143,7 +150,7 @@ def __init__(
                 cost_objective = GenericMCObjective(lambda Y, X: Y.sum(dim=-1))
 
         self.cost_model = cost_model
-        self.cost_objective = cost_objective
+        self.cost_objective: MCAcquisitionObjective = cost_objective
         self._use_mean = use_mean
         self._min_cost = min_cost
 
@@ -153,7 +160,6 @@ def forward(
         deltas: Tensor,
         sampler: Optional[MCSampler] = None,
         X_evaluation_mask: Optional[Tensor] = None,
-        **kwargs: Any,
     ) -> Tensor:
         r"""Evaluate the cost function on the candidates and improvements. Note
         that negative values of `deltas` are instead scaled by the cost, and not

botorch/acquisition/max_value_entropy_search.py

@@ -33,7 +33,7 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from math import log
-from typing import Any, Callable, Optional
+from typing import Callable, Optional
 
 import numpy as np
 import torch
@@ -155,15 +155,14 @@ def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
     # ------- Abstract methods that need to be implemented by subclasses ------- #
 
     @abstractmethod
-    def _compute_information_gain(self, X: Tensor, **kwargs: Any) -> Tensor:
+    def _compute_information_gain(self, X: Tensor) -> Tensor:
         r"""Compute the information gain at the design points `X`.
 
         `num_fantasies = 1` for non-fantasized models.
 
          Args:
             X: A `batch_shape x 1 x d`-dim Tensor of `batch_shape` t-batches
                 with `1` `d`-dim design point each.
-            kwargs: Other keyword arguments used by subclasses.
 
         Returns:
             A `num_fantasies x batch_shape`-dim Tensor of information gains at the
@@ -174,11 +173,12 @@ def _compute_information_gain(self, X: Tensor, **kwargs: Any) -> Tensor:
     @abstractmethod
     def _sample_max_values(
         self, num_samples: int, X_pending: Optional[Tensor] = None
-    ) -> Tensor:
+    ) -> None:
         r"""Draw samples from the posterior over maximum values.
 
         These samples are used to compute Monte Carlo approximations of expectations
-        over the posterior over the function maximum.
+        over the posterior over the function maximum. This function sets
+        `self.posterior_max_values`.
 
         Args:
             num_samples: The number of samples to draw.
@@ -254,11 +254,12 @@ def __init__(
 
     def _sample_max_values(
         self, num_samples: int, X_pending: Optional[Tensor] = None
-    ) -> Tensor:
+    ) -> None:
         r"""Draw samples from the posterior over maximum values on a discrete set.
 
         These samples are used to compute Monte Carlo approximations of expectations
-        over the posterior over the function maximum.
+        over the posterior over the function maximum. This function sets
+        `self.posterior_max_values`.
 
         Args:
             num_samples: The number of samples to draw.

botorch/acquisition/multi_objective/hypervolume_knowledge_gradient.py

@@ -84,7 +84,6 @@ def __init__(
         current_value: Optional[Tensor] = None,
         use_posterior_mean: bool = True,
         cost_aware_utility: Optional[CostAwareUtility] = None,
-        **kwargs: Any,
     ) -> None:
         r"""q-Hypervolume Knowledge Gradient.
 

botorch/acquisition/multi_objective/joint_entropy_search.py

@@ -19,7 +19,7 @@
 
 from abc import abstractmethod
 from math import pi
-from typing import Any, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 from botorch import settings
@@ -50,7 +50,6 @@ def __init__(
         X_pending: Optional[Tensor] = None,
         estimation_type: str = "LB",
         num_samples: int = 64,
-        **kwargs: Any,
     ) -> None:
         r"""Lower bound multi-objective entropy search acquisition function.
 
@@ -283,7 +282,6 @@ def __init__(
         X_pending: Optional[Tensor] = None,
         estimation_type: str = "LB",
         num_samples: int = 64,
-        **kwargs: Any,
     ) -> None:
         r"""Lower bound multi-objective joint entropy search acquisition function.
 

botorch/acquisition/multi_objective/max_value_entropy_search.py

@@ -21,7 +21,7 @@
 
 from math import pi
 
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Callable, Optional, Tuple, Union
 
 import torch
 from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy
@@ -78,7 +78,6 @@ def __init__(
         num_fantasies: int = 16,
         X_pending: Optional[Tensor] = None,
         sampler: Optional[MCSampler] = None,
-        **kwargs: Any,
     ) -> None:
         r"""Multi-objective max-value entropy search acquisition function.
 
@@ -165,7 +164,9 @@ def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
             self._sample_max_values()
 
     def _sample_max_values(self) -> None:
-        r"""Sample max values for MC approximation of the expectation in MES"""
+        """Sample max values for MC approximation of the expectation in MES.
+
+        Sets self.posterior_max_values."""
         with torch.no_grad():
             # num_samples x (num_fantasies) x n_pareto_points x m
             sampled_pfs = self.sample_pareto_frontiers(self.mo_model)
@@ -220,7 +221,6 @@ def __init__(
         X_pending: Optional[Tensor] = None,
         estimation_type: str = "LB",
         num_samples: int = 64,
-        **kwargs: Any,
     ) -> None:
         r"""Lower bound multi-objective max-value entropy search acquisition function.
 

botorch/acquisition/multi_objective/multi_fidelity.py

@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable, List, Optional, Union
+from typing import Callable, List, Optional, Union
 
 import torch
 from botorch.acquisition.cost_aware import InverseCostWeightedUtility
@@ -48,8 +48,7 @@ def __init__(
         constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
         eta: Optional[Union[Tensor, float]] = 1e-3,
         X_pending: Optional[Tensor] = None,
-        cost_call: Callable[Tensor, Tensor] = None,
-        **kwargs: Any,
+        cost_call: Optional[Callable[[Tensor], Tensor]] = None,
     ) -> None:
         r"""MOMF acquisition function supporting m>=2 outcomes.
         The model needs to have train_obj that has a fidelity

botorch/acquisition/multi_objective/objective.py

@@ -28,7 +28,7 @@ class MCMultiOutputObjective(MCAcquisitionObjective):
     _is_mo: bool = True
 
     @abstractmethod
-    def forward(self, samples: Tensor, X: Optional[Tensor] = None, **kwargs) -> Tensor:
+    def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
         r"""Evaluate the multi-output objective on the samples.
 
         Args:

botorch/acquisition/multi_objective/predictive_entropy_search.py

@@ -23,7 +23,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, Tuple
+from typing import Optional, Tuple
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
@@ -107,7 +107,6 @@ def __init__(
         ep_jitter: float = 1e-4,
         test_jitter: float = 1e-4,
         threshold: float = 1e-2,
-        **kwargs: Any,
     ) -> None:
         r"""Multi-objective predictive entropy search acquisition function.
 

botorch/acquisition/multi_step_lookahead.py

@@ -598,7 +598,6 @@ def warmstart_multistep(
     num_restarts: int,
     raw_samples: int,
     full_optimizer: Tensor,
-    **kwargs: Any,
 ) -> Tensor:
     r"""Warm-start initialization for multi-step look-ahead acquisition functions.
 
@@ -614,7 +613,6 @@ def warmstart_multistep(
         full_optimizer: The full tree of optimizers of the previous iteration of shape
             `batch_shape x q' x d`. Typically obtained by passing
             `return_best_only=False` and `return_full_tree=True` into `optimize_acqf`.
-        kwargs: Optimization kwargs.
 
     Returns:
         A `num_restarts x q' x d` tensor for initial points for optimization.

botorch/acquisition/predictive_entropy_search.py

@@ -15,7 +15,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Optional
 
 from botorch.acquisition.multi_objective.predictive_entropy_search import (
     qMultiObjectivePredictiveEntropySearch,
@@ -53,7 +53,6 @@ def __init__(
         ep_jitter: float = 1e-4,
         test_jitter: float = 1e-4,
         threshold: float = 1e-2,
-        **kwargs: Any,
     ) -> None:
         r"""Predictive entropy search acquisition function.
 

botorch/acquisition/preference.py

@@ -28,7 +28,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Optional
 
 import torch
 from botorch.acquisition import AnalyticAcquisitionFunction
@@ -211,7 +211,6 @@ def __init__(
         outcome_model: Optional[DeterministicModel] = None,
         num_samples: Optional[int] = 1024,
         std_noise: Optional[float] = 0.0,
-        **kwargs: Any,
     ) -> None:
         """
         Monte Carlo implementation of Bayesian Active Learning by Disagreement (BALD)

botorch/generation/sampling.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Union
+from typing import Optional, Union
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
@@ -45,14 +45,13 @@ class SamplingStrategy(Module, ABC):
     """
 
     @abstractmethod
-    def forward(self, X: Tensor, num_samples: int = 1, **kwargs: Any) -> Tensor:
+    def forward(self, X: Tensor, num_samples: int = 1) -> Tensor:
         r"""Sample according to the SamplingStrategy.
 
         Args:
             X: A `batch_shape x N x d`-dim Tensor from which to sample (in the `N`
                 dimension).
             num_samples: The number of samples to draw.
-            kwargs: Additional implementation-specific kwargs.
 
         Returns:
             A `batch_shape x num_samples x d`-dim Tensor of samples from `X`, where

botorch/models/approximate_gp.py

@@ -159,7 +159,7 @@ def posterior(
             posterior = self.outcome_transform.untransform_posterior(posterior)
         return posterior
 
-    def forward(self, X, *args, **kwargs) -> MultivariateNormal:
+    def forward(self, X) -> MultivariateNormal:
         if self.training:
             X = self.transform_inputs(X)
         return self.model(X)