code: def test_acquisition_functions(self):
      tkwargs = {"device": self.device, "dtype": torch.double}
      train_X, train_Y, train_Yvar, model = self._get_data_and_model(
          infer_noise=True, **tkwargs
      )
      fit_fully_bayesian_model_nuts(
          model,
          warmup_steps=8,
          num_samples=5,
          thinning=2,
          disable_progbar=True,
      )
      deterministic = GenericDeterministicModel(f=lambda x: x[..., :1])
      # due to ModelList type, setting cache_root=False for all noisy EI variants
      list_gp = ModelListGP(model, model)
      mixed_list = ModelList(deterministic, model)
      simple_sampler = get_sampler(
          posterior=model.posterior(train_X),
          sample_shape=torch.Size([2]),
      )
      list_gp_sampler = get_sampler(
          posterior=list_gp.posterior(train_X),
          sample_shape=torch.Size([2]),
      )
      mixed_list_sampler = get_sampler(
          posterior=mixed_list.posterior(train_X),
          sample_shape=torch.Size([2]),
      )
      acquisition_functions = [
          ExpectedImprovement(model=model, best_f=train_Y.max()),
          ProbabilityOfImprovement(model=model, best_f=train_Y.max()),
          PosteriorMean(model=model),
          UpperConfidenceBound(model=model, beta=4),
          qLogExpectedImprovement(
              model=model, best_f=train_Y.max(), sampler=simple_sampler
          ),
          qExpectedImprovement(
              model=model, best_f=train_Y.max(), sampler=simple_sampler
          ),
          qNoisyExpectedImprovement(
              model=model,
              X_baseline=train_X,
              sampler=simple_sampler,
              cache_root=False,
          ),
          qLogNoisyExpectedImprovement(
              model=model,
              X_baseline=train_X,
              sampler=simple_sampler,
              cache_root=False,
          ),
          qProbabilityOfImprovement(
              model=model, best_f=train_Y.max(), sampler=simple_sampler
          ),
          qSimpleRegret(model=model, sampler=simple_sampler),
          qUpperConfidenceBound(
              model=model, beta=4, sampler=simple_sampler
          ),
          qNoisyExpectedHypervolumeImprovement(
              model=list_gp,
              X_baseline=train_X,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=list_gp_sampler,
              cache_root=False,
          ),
          qLogNoisyExpectedHypervolumeImprovement(
              model=list_gp,
              X_baseline=train_X,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=list_gp_sampler,
              cache_root=False,
          ),
          qExpectedHypervolumeImprovement(
              model=list_gp,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=list_gp_sampler,
              partitioning=NondominatedPartitioning(
                  ref_point=torch.zeros(2, **tkwargs),
                  Y=train_Y.repeat([1, 2]),
              ),
          ),
          qLogExpectedHypervolumeImprovement(
              model=list_gp,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=list_gp_sampler,
              partitioning=NondominatedPartitioning(
                  ref_point=torch.zeros(2, **tkwargs),
                  Y=train_Y.repeat([1, 2]),
              ),
          ),
          # qEHVI/qNEHVI with mixed models
          qNoisyExpectedHypervolumeImprovement(
              model=mixed_list,
              X_baseline=train_X,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=mixed_list_sampler,
              cache_root=False,
          ),
          qLogNoisyExpectedHypervolumeImprovement(
              model=mixed_list,
              X_baseline=train_X,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=mixed_list_sampler,
              cache_root=False,
          ),
          qExpectedHypervolumeImprovement(
              model=mixed_list,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=mixed_list_sampler,
              partitioning=NondominatedPartitioning(
                  ref_point=torch.zeros(2, **tkwargs),
                  Y=train_Y.repeat([1, 2]),
              ),
          ),
          qLogExpectedHypervolumeImprovement(
              model=mixed_list,
              ref_point=torch.zeros(2, **tkwargs),
              sampler=mixed_list_sampler,
              partitioning=NondominatedPartitioning(
                  ref_point=torch.zeros(2, **tkwargs),
                  Y=train_Y.repeat([1, 2]),
              ),
          ),
      ]
  
      for acqf in acquisition_functions:
          for batch_shape in [[5], [6, 5, 2]]:
              test_X = torch.rand(*batch_shape, 1, 4, **tkwargs)
              # Testing that the t_batch_mode_transform works correctly for
              # fully Bayesian models with log-space acquisition functions.
              with patch.object(
                  utils.transforms, "logmeanexp", wraps=logmeanexp
              ) as mock:
                  self.assertEqual(
                      acqf(test_X).shape, torch.Size(batch_shape)
                  )
                  if acqf._log:
                      mock.assert_called_once()
                  else:
                      mock.assert_not_called()
  
      # Test prune_inferior_points
      X_pruned = prune_inferior_points(model=model, X=train_X)
      self.assertTrue(X_pruned.ndim == 2 and X_pruned.shape[-1] == 4)
  
      # Test prune_inferior_points_multi_objective
      for model_list in [
          ModelListGP(model, model),
          ModelList(deterministic, model),
      ]:
          X_pruned = prune_inferior_points_multi_objective(
              model=model_list,
              X=train_X,
              ref_point=torch.zeros(2, **tkwargs),
          )
          self.assertTrue(
              X_pruned.ndim == 2 and X_pruned.shape[-1] == 4
          )
  
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 8 for calling the TestFullyBayesianSingleTaskGP obj's underlying member method code with example usage: obj.test_acquisition_functions() without return types