Don't allow unused keyword arguments in Model.construct_inputs (#2186)

Summary: Pull Request resolved: #2186

Differential Revision: D53086323

Author: esantorella

botorch/models/contextual.py

@@ -47,7 +47,6 @@ def construct_inputs(
         cls,
         training_data: SupervisedDataset,
         decomposition: Dict[str, List[int]],
-        **kwargs: Any,
     ) -> Dict[str, Any]:
         r"""Construct `Model` keyword arguments from a dict of `SupervisedDataset`.
 
@@ -56,7 +55,7 @@ def construct_inputs(
             decomposition: Dictionary of context names and their indexes of the
                 corresponding active context parameters.
         """
-        base_inputs = super().construct_inputs(training_data=training_data, **kwargs)
+        base_inputs = super().construct_inputs(training_data=training_data)
         return {
             **base_inputs,
             "decomposition": decomposition,
@@ -127,7 +126,6 @@ def construct_inputs(
         embs_feature_dict: Optional[Dict] = None,
         embs_dim_list: Optional[List[int]] = None,
         context_weight_dict: Optional[Dict] = None,
-        **kwargs: Any,
     ) -> Dict[str, Any]:
         r"""Construct `Model` keyword arguments from a dict of `SupervisedDataset`.
 
@@ -147,7 +145,7 @@ def construct_inputs(
                 dimension is set to 1 for each categorical variable.
             context_weight_dict: Known population weights of each context.
         """
-        base_inputs = super().construct_inputs(training_data=training_data, **kwargs)
+        base_inputs = super().construct_inputs(training_data=training_data)
         index_decomp = {
             c: [training_data.feature_names.index(i) for i in v]
             for c, v in decomposition.items()

botorch/models/fully_bayesian_multitask.py

@@ -388,18 +388,28 @@ def construct_inputs(
         cls,
         training_data: Union[SupervisedDataset, MultiTaskDataset],
         task_feature: int,
+        output_tasks: Optional[List[int]] = None,
+        prior_config: Optional[dict] = None,
         rank: Optional[int] = None,
-        **kwargs: Any,
     ) -> Dict[str, Any]:
         r"""Construct `Model` keyword arguments from a dataset and other args.
 
         Args:
             training_data: A `SupervisedDataset` or a `MultiTaskDataset`.
             task_feature: Column index of embedded task indicator features.
+            output_tasks: A list of task indices for which to compute model
+                outputs for. If omitted, return outputs for all task indices.
+            prior_config: Configuration for inter-task covariance prior.
+                Should only be used if `task_covar_prior` is not passed directly. Must
+                contain `use_LKJ_prior` indicator and should contain float value `eta`.
             rank: The rank of the cross-task covariance matrix.
         """
         inputs = super().construct_inputs(
-            training_data=training_data, task_feature=task_feature, rank=rank, **kwargs
+            training_data=training_data,
+            task_feature=task_feature,
+            rank=rank,
+            output_tasks=output_tasks,
+            prior_config=prior_config,
         )
         inputs.pop("task_covar_prior")
         if "train_Yvar" not in inputs:

botorch/models/gp_regression_fidelity.py

@@ -167,15 +167,14 @@ def construct_inputs(
         cls,
         training_data: SupervisedDataset,
         fidelity_features: List[int],
-        **kwargs,
     ) -> Dict[str, Any]:
         r"""Construct `Model` keyword arguments from a dict of `SupervisedDataset`.
 
         Args:
             training_data: Dictionary of `SupervisedDataset`.
             fidelity_features: Index of fidelity parameter as input columns.
         """
-        inputs = super().construct_inputs(training_data=training_data, **kwargs)
+        inputs = super().construct_inputs(training_data=training_data)
         inputs["data_fidelities"] = fidelity_features
         return inputs
 

botorch/models/gp_regression_mixed.py

@@ -187,7 +187,6 @@ def construct_inputs(
         training_data: SupervisedDataset,
         categorical_features: List[int],
         likelihood: Optional[Likelihood] = None,
-        **kwargs: Any,
     ) -> Dict[str, Any]:
         r"""Construct `Model` keyword arguments from a dict of `SupervisedDataset`.
 
@@ -196,7 +195,7 @@ def construct_inputs(
             categorical_features: Column indices of categorical features.
             likelihood: Optional likelihood used to constuct the model.
         """
-        base_inputs = super().construct_inputs(training_data=training_data, **kwargs)
+        base_inputs = super().construct_inputs(training_data=training_data)
         return {
             **base_inputs,
             "cat_dims": categorical_features,

botorch/models/model.py

@@ -187,15 +187,13 @@ def condition_on_observations(self, X: Tensor, Y: Tensor, **kwargs: Any) -> Mode
     def construct_inputs(
         cls,
         training_data: SupervisedDataset,
-        **kwargs: Any,
     ) -> Dict[str, Union[BotorchContainer, Tensor]]:
         """
         Construct `Model` keyword arguments from a `SupervisedDataset`.
 
         Args:
             training_data: A `SupervisedDataset`, with attributes `train_X`,
                 `train_Y`, and, optionally, `train_Yvar`.
-            kwargs: Ignored.
 
         Returns:
             A dict of keyword arguments that can be used to initialize a `Model`,

botorch/models/multitask.py

@@ -277,7 +277,6 @@ def construct_inputs(
         task_covar_prior: Optional[Prior] = None,
         prior_config: Optional[dict] = None,
         rank: Optional[int] = None,
-        **kwargs,
     ) -> Dict[str, Any]:
         r"""Construct `Model` keyword arguments from a dataset and other args.
 
@@ -310,9 +309,8 @@ def construct_inputs(
                 raise ValueError(f"eta must be a real number, your eta was {eta}.")
             task_covar_prior = LKJCovariancePrior(num_tasks, eta, sd_prior)
 
-        base_inputs = super().construct_inputs(
-            training_data=training_data, task_feature=task_feature, **kwargs
-        )
+        # Call Model.construct_inputs to parse training data
+        base_inputs = super().construct_inputs(training_data=training_data)
         return {
             **base_inputs,
             "task_feature": task_feature,

botorch/models/pairwise_gp.py

@@ -788,15 +788,13 @@ def batch_shape(self) -> torch.Size:
     def construct_inputs(
         cls,
         training_data: SupervisedDataset,
-        **kwargs: Any,
     ) -> Dict[str, Tensor]:
         r"""
         Construct `Model` keyword arguments from a `RankingDataset`.
 
         Args:
             training_data: A `RankingDataset`, with attributes `train_X`,
                 `train_Y`, and, optionally, `train_Yvar`.
-            kwargs: Ignored.
 
         Returns:
             A dict of keyword arguments that can be used to initialize a