class_docstr: A fully Bayesian multi-task GP model with the SAAS prior.
  
  This model assumes that the inputs have been normalized to [0, 1]^d and that the
  output has been stratified standardized to have zero mean and unit variance for
  each task.The SAAS model [Eriksson2021saasbo]_ with a Matern-5/2 is used as data
  kernel by default.
  
  You are expected to use `fit_fully_bayesian_model_nuts` to fit this model as it
  isn't compatible with `fit_gpytorch_model`.
  
  Example:
      >>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
      >>> i1, i2 = torch.zeros(10, 1), torch.ones(20, 1)
      >>> train_X = torch.cat([
      >>>     torch.cat([X1, i1], -1), torch.cat([X2, i2], -1),
      >>> ])
      >>> train_Y = torch.cat(f1(X1), f2(X2)).unsqueeze(-1)
      >>> train_Yvar = 0.01 * torch.ones_like(train_Y)
      >>> mtsaas_gp = SaasFullyBayesianMultiTaskGP(
      >>>     train_X, train_Y, train_Yvar, task_feature=-1,
      >>> )
      >>> fit_fully_bayesian_model_nuts(mtsaas_gp)
      >>> posterior = mtsaas_gp.posterior(test_X)
docstr: Initialize the fully Bayesian multi-task GP model.
  
  Args:
      train_X: Training inputs (n x (d + 1))
      train_Y: Training targets (n x 1)
      train_Yvar: Observed noise variance (n x 1). If None, we infer the noise.
          Note that the inferred noise is common across all tasks.
      task_feature: The index of the task feature (`-d <= task_feature <= d`).
      output_tasks: A list of task indices for which to compute model
          outputs for. If omitted, return outputs for all task indices.
      rank: The num of learned task embeddings to be used in the task kernel.
          If omitted, use a full rank (i.e. number of tasks) kernel.
      outcome_transform: An outcome transform that is applied to the
          training data during instantiation and to the posterior during
          inference (that is, the `Posterior` obtained by calling
          `.posterior` on the model will be on the original scale).
      input_transform: An input transform that is applied to the inputs `X`
          in the model's forward pass.
      pyro_model: Optional `PyroModel`, defaults to `MultitaskSaasPyroModel`.
code: def __init__(
      self,
      train_X: Tensor,
      train_Y: Tensor,
      task_feature: int,
      train_Yvar: Optional[Tensor] = None,
      output_tasks: Optional[List[int]] = None,
      rank: Optional[int] = None,
      outcome_transform: Optional[OutcomeTransform] = None,
      input_transform: Optional[InputTransform] = None,
      pyro_model: Optional[PyroModel] = None,
  ) -> None:
      r"""Initialize the fully Bayesian multi-task GP model.
  
      Args:
          train_X: Training inputs (n x (d + 1))
          train_Y: Training targets (n x 1)
          train_Yvar: Observed noise variance (n x 1). If None, we infer the noise.
              Note that the inferred noise is common across all tasks.
          task_feature: The index of the task feature (`-d <= task_feature <= d`).
          output_tasks: A list of task indices for which to compute model
              outputs for. If omitted, return outputs for all task indices.
          rank: The num of learned task embeddings to be used in the task kernel.
              If omitted, use a full rank (i.e. number of tasks) kernel.
          outcome_transform: An outcome transform that is applied to the
              training data during instantiation and to the posterior during
              inference (that is, the `Posterior` obtained by calling
              `.posterior` on the model will be on the original scale).
          input_transform: An input transform that is applied to the inputs `X`
              in the model's forward pass.
          pyro_model: Optional `PyroModel`, defaults to `MultitaskSaasPyroModel`.
      """
      if not (
          train_X.ndim == train_Y.ndim == 2
          and len(train_X) == len(train_Y)
          and train_Y.shape[-1] == 1
      ):
          raise ValueError(
              "Expected train_X to have shape n x d and train_Y to have shape n x 1"
          )
      if train_Yvar is not None and train_Y.shape != train_Yvar.shape:
          raise ValueError(
              "Expected train_Yvar to be None or have the same shape as train_Y"
          )
      with torch.no_grad():
          transformed_X = self.transform_inputs(
              X=train_X, input_transform=input_transform
          )
      if outcome_transform is not None:
          train_Y, train_Yvar = outcome_transform(train_Y, train_Yvar)
      if train_Yvar is not None:  # Clamp after transforming
          train_Yvar = train_Yvar.clamp(MIN_INFERRED_NOISE_LEVEL)
  
      super().__init__(
          train_X=train_X,
          train_Y=train_Y,
          train_Yvar=train_Yvar,
          task_feature=task_feature,
          output_tasks=output_tasks,
          rank=rank,
      )
      self.to(train_X)
  
      self.mean_module = None
      self.covar_module = None
      self.likelihood = None
      self.task_covar_module = None
      self.register_buffer("latent_features", None)
      if pyro_model is None:
          pyro_model = MultitaskSaasPyroModel()
      pyro_model.set_inputs(
          train_X=transformed_X,
          train_Y=train_Y,
          train_Yvar=train_Yvar,
          task_feature=task_feature,
          task_rank=self._rank,
      )
      self.pyro_model = pyro_model
      if outcome_transform is not None:
          self.outcome_transform = outcome_transform
      if input_transform is not None:
          self.input_transform = input_transform
  
imports: from typing import Any, Dict, List, Mapping, NoReturn, Optional, Tuple, Union
  import pyro
  import torch
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.models.fully_bayesian import matern52_kernel, MIN_INFERRED_NOISE_LEVEL, PyroModel, reshape_and_detach, SaasPyroModel
  from botorch.models.multitask import MultiTaskGP
  from botorch.models.transforms.input import InputTransform
  from botorch.models.transforms.outcome import OutcomeTransform
  from botorch.posteriors.fully_bayesian import GaussianMixturePosterior, MCMC_DIM
  from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
  from gpytorch.distributions.multivariate_normal import MultivariateNormal
  from gpytorch.kernels import MaternKernel
  from gpytorch.kernels.kernel import Kernel
  from gpytorch.likelihoods.likelihood import Likelihood
  from gpytorch.means.mean import Mean
  from torch import Tensor
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 1 to create a new SaasFullyBayesianMultiTaskGP object from inherited base classes: MultiTaskGP with example: obj = SaasFullyBayesianMultiTaskGP(train_X, train_Y, task_feature, train_Yvar, output_tasks, rank, outcome_transform, input_transform, pyro_model)