class_docstr: Does Bayesian optimization with LCE-M GP.
class_docstr_tok: ["Does","Bayesian","optimization","with","LCE-M","GP","."]
docstr: Get a fitted multi-task contextual GP model for each outcome.
  Args:
      Xs: List of X data, one tensor per outcome.
      Ys: List of Y data, one tensor per outcome.
      Yvars:List of Noise variance of Yvar data, one tensor per outcome.
      task_features: List of columns of X that are tasks.
  Returns: ModeListGP that each model is a fitted LCEM GP model.
docstr_tok: ["Get","a","fitted","multi-task","contextual","GP","model","for","each","outcome",".","Args",":","Xs",":","List","of","X","data",",","one","tensor","per","outcome",".","Ys",":","List","of","Y","data",",","one","tensor","per","outcome",".","Yvars",":","List","of","Noise","variance","of","Yvar","data",",","one","tensor","per","outcome",".","task_features",":","List","of","columns","of","X","that","are","tasks",".","Returns",":","ModeListGP","that","each","model","is","a","fitted","LCEM","GP","model","."]
code: def get_and_fit_model(
      self,
      Xs: List[Tensor],
      Ys: List[Tensor],
      Yvars: List[Tensor],
      task_features: List[int],
      fidelity_features: List[int],
      metric_names: List[str],
      state_dict: Optional[Dict[str, Tensor]] = None,
      fidelity_model_id: Optional[int] = None,
      **kwargs: Any,
  ) -> ModelListGP:
      """Get a fitted multi-task contextual GP model for each outcome.
      Args:
          Xs: List of X data, one tensor per outcome.
          Ys: List of Y data, one tensor per outcome.
          Yvars:List of Noise variance of Yvar data, one tensor per outcome.
          task_features: List of columns of X that are tasks.
      Returns: ModeListGP that each model is a fitted LCEM GP model.
      """
  
      if len(task_features) == 1:
          task_feature = task_features[0]
      elif len(task_features) > 1:
          raise NotImplementedError(
              f"LCEMBO only supports 1 task feature (got {task_features})"
          )
      else:
          raise ValueError(
              "LCEMBO requires context input as task features"
          )
  
      models = []
      for i, X in enumerate(Xs):
          # validate input Yvars
          Yvar = Yvars[i].clamp_min_(MIN_OBSERVED_NOISE_LEVEL)
          is_nan = torch.isnan(Yvar)
          all_nan_Yvar = torch.all(is_nan)
          if all_nan_Yvar:
              gp_m = LCEMGP(
                  train_X=X,
                  train_Y=Ys[i],
                  task_feature=task_feature,
                  context_cat_feature=self.context_cat_feature,
                  context_emb_feature=self.context_emb_feature,
                  embs_dim_list=self.embs_dim_list,
              )
          else:
              gp_m = LCEMGP(
                  train_X=X,
                  train_Y=Ys[i],
                  train_Yvar=Yvar,
                  task_feature=task_feature,
                  context_cat_feature=self.context_cat_feature,
                  context_emb_feature=self.context_emb_feature,
                  embs_dim_list=self.embs_dim_list,
              )
          models.append(gp_m)
      # Use a ModelListGP
      model = ModelListGP(*models)
      model.to(Xs[0])
      mll = SumMarginalLogLikelihood(model.likelihood, model)
      fit_gpytorch_mll(mll)
      return model
  
file_path: Ax/ax/models/torch/cbo_lcem.py
imports: from typing import Any, Dict, List, Optional
  import torch
  from ax.models.torch.botorch import BotorchModel
  from botorch.fit import fit_gpytorch_mll
  from botorch.models.contextual_multioutput import LCEMGP
  from botorch.models.model_list_gp_regression import ModelListGP
  from gpytorch.mlls.sum_marginal_log_likelihood import SumMarginalLogLikelihood
  from torch import Tensor
label_desc: Use image node_id 2 for calling the LCEMBO obj's underlying member method code with example usage: obj.get_and_fit_model(Xs, Ys, Yvars, task_features, fidelity_features, metric_names, state_dict, fidelity_model_id) and returns: model