add output_tasks to MTGP in MBM (#2241)

Summary:

see title. output_tasks were not set for MTGP in MBM.

Differential Revision: D54453253

ax/models/torch/botorch_modular/surrogate.py

@@ -49,6 +49,7 @@
 )
 from botorch.models.model import Model
 from botorch.models.model_list_gp_regression import ModelListGP
+from botorch.models.multitask import MultiTaskGP
 from botorch.models.pairwise_gp import PairwiseGP
 from botorch.models.transforms.input import (
  ChainedInputTransform,
@@ -59,6 +60,7 @@
 from botorch.utils.containers import SliceContainer
 from botorch.utils.datasets import RankingDataset, SupervisedDataset
 from botorch.utils.dispatcher import Dispatcher
+from botorch_fb.models.heterogeneous_mtgp import HeterogeneousMTGP
 from gpytorch.kernels import Kernel
 from gpytorch.likelihoods.likelihood import Likelihood
 from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
@@ -776,6 +778,10 @@ def outcomes(self, value: List[str]) -> None:
 )
 
 
+# Note: although `HeterogeneousMTGP`` is a subclass of `MultiTaskGP`,
+# it takes different inputs than a `MultiTaskGP`` and it's
+# datasets can be heterogeneous, which means `dataset.X`
+# will fail since it tries to concatenate heterogeneous datasets.
 @submodel_input_constructor.register(Model)
 def _submodel_input_constructor_base(
  botorch_model_class: Type[Model],
@@ -841,3 +847,26 @@ def _submodel_input_constructor_base(
  botorch_model_class_args=botorch_model_class_args,
  )
  return formatted_model_inputs
+
+
+@submodel_input_constructor.register(MultiTaskGP)
+def _submodel_input_constructor_mtgp(
+ botorch_model_class: Type[Model],
+ dataset: SupervisedDataset,
+ search_space_digest: SearchSpaceDigest,
+ surrogate: Surrogate,
+) -> Dict[str, Any]:
+ if len(dataset.outcome_names) > 1:
+ raise NotImplementedError("Multi-output Multi-task GPs are not yet supported.")
+ formatted_model_inputs = _submodel_input_constructor_base(
+ botorch_model_class=botorch_model_class,
+ dataset=dataset,
+ search_space_digest=search_space_digest,
+ surrogate=surrogate,
+ )
+ # specify output tasks so that model.num_outputs = 1
+ # since the model only models a single outcome
+ formatted_model_inputs["output_tasks"] = dataset.X[
+ :1, formatted_model_inputs["task_feature"]
+ ].tolist()
+ return formatted_model_inputs