code: def test_condition_on_observation(self):
      # The following conditioned data shapes should work (output describes):
      # training data shape after cond(batch shape in output is req. in gpytorch)
      # X: num_models x n x d, Y: num_models x n x d --> num_models x n x d
      # X: n x d, Y: n x d --> num_models x n x d
      # X: n x d, Y: num_models x n x d --> num_models x n x d
      num_models = 3
      num_cond = 2
      for infer_noise, dtype in itertools.product(
          (True, False), (torch.float, torch.double)
      ):
          tkwargs = {"device": self.device, "dtype": dtype}
          (
              train_X,
              train_Y,
              train_Yvar,
              test_X,
          ) = self._get_unnormalized_data(
              infer_noise=infer_noise, **tkwargs
          )
          num_train, num_dims = train_X.shape
          # condition on different observations per model to obtain num_models sets
          # of training data
          (
              cond_X,
              cond_Y,
              cond_Yvar,
          ) = self._get_unnormalized_condition_data(
              num_models=num_models,
              num_cond=num_cond,
              infer_noise=infer_noise,
              **tkwargs,
          )
          model = SaasFullyBayesianSingleTaskGP(
              train_X=train_X,
              train_Y=train_Y,
              train_Yvar=train_Yvar,
          )
          mcmc_samples = self._get_mcmc_samples(
              num_samples=num_models,
              dim=train_X.shape[-1],
              infer_noise=infer_noise,
              **tkwargs,
          )
          model.load_mcmc_samples(mcmc_samples)
  
          # need to forward pass before conditioning
          model.posterior(train_X)
          cond_model = model.condition_on_observations(
              cond_X, cond_Y, noise=cond_Yvar
          )
          posterior = cond_model.posterior(test_X)
          self.assertEqual(
              posterior.mean.shape,
              torch.Size([num_models, len(test_X), 1]),
          )
  
          # since the data is not equal for the conditioned points, a batch size
          # is added to the training data
          self.assertEqual(
              cond_model.train_inputs[0].shape,
              torch.Size([num_models, num_train + num_cond, num_dims]),
          )
  
          # the batch shape of the condition model is added during conditioning
          self.assertEqual(
              cond_model.batch_shape, torch.Size([num_models])
          )
  
          # condition on identical sets of data (i.e. one set) for all models
          # i.e, with no batch shape. This infers the batch shape.
          cond_X_nobatch, cond_Y_nobatch = cond_X[0], cond_Y[0]
          model = SaasFullyBayesianSingleTaskGP(
              train_X=train_X,
              train_Y=train_Y,
              train_Yvar=train_Yvar,
          )
          mcmc_samples = self._get_mcmc_samples(
              num_samples=num_models,
              dim=train_X.shape[-1],
              infer_noise=infer_noise,
              **tkwargs,
          )
          model.load_mcmc_samples(mcmc_samples)
  
          # conditioning without a batch size - the resulting conditioned model
          # will still have a batch size
          model.posterior(train_X)
          cond_model = model.condition_on_observations(
              cond_X_nobatch, cond_Y_nobatch, noise=cond_Yvar
          )
          self.assertEqual(
              cond_model.train_inputs[0].shape,
              torch.Size([num_models, num_train + num_cond, num_dims]),
          )
  
          # test repeated conditining
          repeat_cond_X = cond_X + 5
          repeat_cond_model = cond_model.condition_on_observations(
              repeat_cond_X, cond_Y, noise=cond_Yvar
          )
          self.assertEqual(
              repeat_cond_model.train_inputs[0].shape,
              torch.Size(
                  [num_models, num_train + 2 * num_cond, num_dims]
              ),
          )
  
          # test repeated conditioning without a batch size
          repeat_cond_X_nobatch = cond_X_nobatch + 10
          repeat_cond_model2 = (
              repeat_cond_model.condition_on_observations(
                  repeat_cond_X_nobatch, cond_Y_nobatch, noise=cond_Yvar
              )
          )
          self.assertEqual(
              repeat_cond_model2.train_inputs[0].shape,
              torch.Size(
                  [num_models, num_train + 3 * num_cond, num_dims]
              ),
          )
  
imports: import itertools
  from unittest import mock
  from unittest.mock import patch
  import pyro
  import torch
  from botorch import fit_fully_bayesian_model_nuts, utils
  from botorch.acquisition.analytic import ExpectedImprovement, PosteriorMean, ProbabilityOfImprovement, UpperConfidenceBound
  from botorch.acquisition.logei import qLogExpectedImprovement, qLogNoisyExpectedImprovement
  from botorch.acquisition.monte_carlo import qExpectedImprovement, qNoisyExpectedImprovement, qProbabilityOfImprovement, qSimpleRegret, qUpperConfidenceBound
  from botorch.acquisition.multi_objective import prune_inferior_points_multi_objective, qExpectedHypervolumeImprovement, qNoisyExpectedHypervolumeImprovement
  from botorch.acquisition.multi_objective.logei import qLogExpectedHypervolumeImprovement, qLogNoisyExpectedHypervolumeImprovement
  from botorch.acquisition.utils import prune_inferior_points
  from botorch.models import ModelList, ModelListGP
  from botorch.models.deterministic import GenericDeterministicModel
  from botorch.models.fully_bayesian import MCMC_DIM, MIN_INFERRED_NOISE_LEVEL, PyroModel, SaasFullyBayesianSingleTaskGP, SaasPyroModel
  from botorch.models.transforms import Normalize, Standardize
  from botorch.posteriors.fully_bayesian import batched_bisect, GaussianMixturePosterior
  from botorch.sampling.get_sampler import get_sampler
  from botorch.utils.datasets import SupervisedDataset
  from botorch.utils.multi_objective.box_decompositions.non_dominated import NondominatedPartitioning
  from botorch.utils.safe_math import logmeanexp
  from botorch.utils.testing import BotorchTestCase
  from gpytorch.distributions import MultivariateNormal
  from gpytorch.kernels import MaternKernel, ScaleKernel
  from gpytorch.likelihoods import FixedNoiseGaussianLikelihood, GaussianLikelihood
  from gpytorch.means import ConstantMean
  from linear_operator.operators import to_linear_operator
  from pyro.ops.integrator import _EXCEPTION_HANDLERS, potential_grad, register_exception_handler
label_desc: Use image node_id 12 for calling the TestFullyBayesianSingleTaskGP obj's underlying member method code with example usage: obj.test_condition_on_observation() without return types