class_docstr: Implementation of the multi-task sparse axis-aligned subspace priors (SAAS) model.
  
  The multi-task model uses an ICM kernel. The data kernel is same as the single task
  SAAS model in order to handle high-dimensional parameter spaces. The task kernel
  is a Matern-5/2 kernel using learned task embeddings as the input.
docstr: Set the training data.
  
  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`).
      task_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.
code: def set_inputs(
      self,
      train_X: Tensor,
      train_Y: Tensor,
      train_Yvar: Optional[Tensor],
      task_feature: int,
      task_rank: Optional[int] = None,
  ):
      """Set the training data.
  
      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`).
          task_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.
      """
      super().set_inputs(train_X, train_Y, train_Yvar)
      # obtain a list of task indicies
      all_tasks = (
          train_X[:, task_feature]
          .unique()
          .to(dtype=torch.long)
          .tolist()
      )
      self.task_feature = task_feature
      self.num_tasks = len(all_tasks)
      self.task_rank = task_rank or self.num_tasks
      # assume there is one column for task feature
      self.ard_num_dims = self.train_X.shape[-1] - 1
  
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 for calling the MultitaskSaasPyroModel obj's underlying member method code with example usage: obj.set_inputs(train_X, train_Y, train_Yvar, task_feature, task_rank) without return types