Don't allow unused **kwargs in input_constructors except for a defined set of exceptions (meta-pytorch#1872)

Summary:
X-link: facebook/Ax#1772

Pull Request resolved: meta-pytorch#1872

[x] Remove unused arguments from input constructors and related functions. The idea is especially not to let unused keyword arguments disappear into `**kwargs` and be silently ignored
[x] add arguments to some input constructors so they don't need any `**kwargs`
[x] Add a decorator that ensures that each input constructor can accept a certain set of keyword arguments, even if those are not used are the constructor, while still erroring on
[ ] Prevent arguments from having different defaults in the input constructors as in acquisition functions

Reviewed By: lena-kashtelyan

Differential Revision: D46519588

fbshipit-source-id: 862fab3fca2460e04462b4b408b6bc3431baa0ff

Author: esantorella

botorch/acquisition/input_constructors.py

@@ -25,7 +25,7 @@
 import warnings
 from math import log, pi
 
-from typing import Any, Optional
+from typing import Optional
 
 import torch
 from botorch import settings
@@ -78,7 +78,6 @@ def __init__(
         estimation_type: str = "LB",
         maximize: bool = True,
         num_samples: int = 64,
-        **kwargs: Any,
     ) -> None:
         r"""Joint entropy search acquisition function.
 

botorch/acquisition/joint_entropy_search.py

@@ -74,7 +74,6 @@ def __init__(
         inner_sampler: Optional[MCSampler] = None,
         X_pending: Optional[Tensor] = None,
         current_value: Optional[Tensor] = None,
-        **kwargs: Any,
     ) -> None:
         r"""q-Knowledge Gradient (one-shot optimization).
 
@@ -330,7 +329,6 @@ def __init__(
         expand: Callable[[Tensor], Tensor] = lambda X: X,
         valfunc_cls: Optional[Type[AcquisitionFunction]] = None,
         valfunc_argfac: Optional[Callable[[Model], Dict[str, Any]]] = None,
-        **kwargs: Any,
     ) -> None:
         r"""Multi-Fidelity q-Knowledge Gradient (one-shot optimization).
 

botorch/acquisition/knowledge_gradient.py

@@ -326,7 +326,6 @@ def __init__(
         maximize: bool = True,
         X_pending: Optional[Tensor] = None,
         train_inputs: Optional[Tensor] = None,
-        **kwargs: Any,
     ) -> None:
         r"""Single-outcome max-value entropy search acquisition function.
 
@@ -697,7 +696,6 @@ def __init__(
         cost_aware_utility: Optional[CostAwareUtility] = None,
         project: Callable[[Tensor], Tensor] = lambda X: X,
         expand: Callable[[Tensor], Tensor] = lambda X: X,
-        **kwargs: Any,
     ) -> None:
         r"""Single-outcome max-value entropy search acquisition function.
 

botorch/acquisition/max_value_entropy_search.py

@@ -24,7 +24,7 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from functools import partial
-from typing import Any, Callable, List, Optional, Protocol, Tuple, Union
+from typing import Callable, List, Optional, Protocol, Tuple, Union
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
@@ -351,7 +351,6 @@ def __init__(
         X_pending: Optional[Tensor] = None,
         constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
         eta: Union[Tensor, float] = 1e-3,
-        **kwargs: Any,
     ) -> None:
         r"""q-Expected Improvement.
 
@@ -434,7 +433,7 @@ def __init__(
         cache_root: bool = True,
         constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
         eta: Union[Tensor, float] = 1e-3,
-        **kwargs: Any,
+        marginalize_dim: Optional[int] = None,
     ) -> None:
         r"""q-Noisy Expected Improvement.
 
@@ -469,6 +468,7 @@ def __init__(
             eta: Temperature parameter(s) governing the smoothness of the sigmoid
                 approximation to the constraint indicators. For more details, on this
                 parameter, see the docs of `compute_smoothed_feasibility_indicator`.
+            marginalize_dim: The dimension to marginalize over.
 
         TODO: similar to qNEHVI, when we are using sequential greedy candidate
         selection, we could incorporate pending points X_baseline and compute
@@ -491,7 +491,7 @@ def __init__(
                 X=X_baseline,
                 objective=objective,
                 posterior_transform=posterior_transform,
-                marginalize_dim=kwargs.get("marginalize_dim"),
+                marginalize_dim=marginalize_dim,
             )
         self.register_buffer("X_baseline", X_baseline)
         # registering buffers for _get_samples_and_objectives in the next `if` block

botorch/acquisition/monte_carlo.py

@@ -27,7 +27,7 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from itertools import combinations
-from typing import Any, Callable, List, Optional, Union
+from typing import Callable, List, Optional, Union
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
@@ -373,7 +373,7 @@ def __init__(
         max_iep: int = 0,
         incremental_nehvi: bool = True,
         cache_root: bool = True,
-        **kwargs: Any,
+        marginalize_dim: Optional[int] = None,
     ) -> None:
         r"""q-Noisy Expected Hypervolume Improvement supporting m>=2 outcomes.
 
@@ -466,7 +466,7 @@ def __init__(
                 objective=objective,
                 constraints=constraints,
                 ref_point=ref_point,
-                marginalize_dim=kwargs.get("marginalize_dim"),
+                marginalize_dim=marginalize_dim,
             )
         self.register_buffer("ref_point", ref_point)
         self.alpha = alpha