code: def choose_model_class(
      datasets: List[SupervisedDataset],
      search_space_digest: SearchSpaceDigest,
  ) -> Type[Model]:
      """Chooses a BoTorch `Model` using the given data (currently just Yvars)
      and its properties (information about task and fidelity features).
  
      Args:
          Yvars: List of tensors, each representing observation noise for a
              given outcome, where outcomes are in the same order as in Xs.
          task_features: List of columns of X that are tasks.
          fidelity_features: List of columns of X that are fidelity parameters.
  
      Returns:
          A BoTorch `Model` class.
      """
      if len(search_space_digest.fidelity_features) > 1:
          raise NotImplementedError(
              "Only a single fidelity feature supported "
              f"(got: {search_space_digest.fidelity_features})."
          )
      if len(search_space_digest.task_features) > 1:
          raise NotImplementedError(
              f"Only a single task feature supported "
              f"(got: {search_space_digest.task_features})."
          )
      if (
          search_space_digest.task_features
          and search_space_digest.fidelity_features
      ):
          raise NotImplementedError(
              "Multi-task multi-fidelity optimization not yet supported."
          )
  
      is_fixed_noise = [ds.Yvar is not None for ds in datasets]
      all_inferred = not any(is_fixed_noise)
      if not all_inferred and not all(is_fixed_noise):
          raise ValueError(
              "Mix of known and unknown variances indicates valuation function "
              "errors. Variances should all be specified, or none should be."
          )
  
      # Multi-task case (when `task_features` is specified).
      if search_space_digest.task_features:
          model_class = MultiTaskGP
  
      # Single-task multi-fidelity cases.
      elif search_space_digest.fidelity_features:
          model_class = SingleTaskMultiFidelityGP
  
      # Mixed optimization case. Note that presence of categorical
      # features in search space digest indicates that downstream in the
      # stack we chose not to perform continuous relaxation on those
      # features.
      elif search_space_digest.categorical_features:
          model_class = MixedSingleTaskGP
  
      # Single-task single-fidelity cases.
      else:
          model_class = SingleTaskGP
  
      logger.debug(f"Chose BoTorch model class: {model_class}.")
      return model_class
  
imports: import warnings
  from logging import Logger
  from typing import Any, Callable, Dict, List, Optional, OrderedDict, Tuple, Type
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.exceptions.core import AxError, AxWarning, UnsupportedError
  from ax.models.types import TConfig
  from ax.utils.common.constants import Keys
  from ax.utils.common.logger import get_logger
  from ax.utils.common.typeutils import checked_cast, not_none
  from botorch.acquisition.acquisition import AcquisitionFunction
  from botorch.acquisition.logei import qLogNoisyExpectedImprovement
  from botorch.acquisition.multi_objective.logei import qLogNoisyExpectedHypervolumeImprovement
  from botorch.fit import fit_fully_bayesian_model_nuts, fit_gpytorch_mll
  from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
  from botorch.models.gp_regression import SingleTaskGP
  from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
  from botorch.models.gp_regression_mixed import MixedSingleTaskGP
  from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel, GPyTorchModel
  from botorch.models.model import Model, ModelList
  from botorch.models.multitask import MultiTaskGP
  from botorch.models.pairwise_gp import PairwiseGP
  from botorch.utils.datasets import SupervisedDataset
  from botorch.utils.transforms import is_fully_bayesian
  from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
  from torch import Tensor
label_desc: Use image node_id 2 for calling a global function with example usage: choose_model_class(datasets, search_space_digest) and returns: model_class