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: Sample from the SAAS model.
  
  This samples the mean, noise variance, outputscale, and lengthscales according
  to the SAAS prior.
code: def sample(self) -> None:
      r"""Sample from the SAAS model.
  
      This samples the mean, noise variance, outputscale, and lengthscales according
      to the SAAS prior.
      """
      tkwargs = {
          "dtype": self.train_X.dtype,
          "device": self.train_X.device,
      }
      base_idxr = torch.arange(
          self.ard_num_dims, **{"device": tkwargs["device"]}
      )
      base_idxr[self.task_feature :] += 1  # exclude task feature
      task_indices = self.train_X[..., self.task_feature].to(
          device=tkwargs["device"], dtype=torch.long
      )
  
      outputscale = self.sample_outputscale(
          concentration=2.0, rate=0.15, **tkwargs
      )
      mean = self.sample_mean(**tkwargs)
      noise = self.sample_noise(**tkwargs)
  
      lengthscale = self.sample_lengthscale(
          dim=self.ard_num_dims, **tkwargs
      )
      K = matern52_kernel(
          X=self.train_X[..., base_idxr], lengthscale=lengthscale
      )
  
      # compute task covar matrix
      task_latent_features = self.sample_latent_features(**tkwargs)[
          task_indices
      ]
      task_lengthscale = self.sample_task_lengthscale(**tkwargs)
      task_covar = matern52_kernel(
          X=task_latent_features, lengthscale=task_lengthscale
      )
      K = K.mul(task_covar)
      K = outputscale * K + noise * torch.eye(
          self.train_X.shape[0], **tkwargs
      )
      pyro.sample(
          "Y",
          pyro.distributions.MultivariateNormal(
              loc=mean.view(-1).expand(self.train_X.shape[0]),
              covariance_matrix=K,
          ),
          obs=self.train_Y.squeeze(-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 2 for calling the MultitaskSaasPyroModel obj's underlying member method code with example usage: obj.sample() without return types