code: def test_get_X_baseline(self):
      tkwargs = {"device": self.device}
      for dtype in (torch.float, torch.double):
          tkwargs["dtype"] = dtype
          X_train = torch.rand(20, 2, **tkwargs)
          model = MockModel(
              MockPosterior(
                  mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)
              )
          )
          # test NEI with X_baseline
          acqf = qNoisyExpectedImprovement(
              model,
              X_baseline=X_train[:2],
              prune_baseline=False,
              cache_root=False,
          )
          X = get_X_baseline(acq_function=acqf)
          self.assertTrue(torch.equal(X, acqf.X_baseline))
          # test EI without X_baseline
          acqf = qExpectedImprovement(model, best_f=0.0)
  
          with warnings.catch_warnings(
              record=True
          ) as w, settings.debug(True):
              X_rnd = get_X_baseline(
                  acq_function=acqf,
              )
              self.assertEqual(len(w), 1)
              self.assertTrue(
                  issubclass(w[-1].category, BotorchWarning)
              )
              self.assertIsNone(X_rnd)
  
          # set train inputs
          model.train_inputs = (X_train,)
          X = get_X_baseline(
              acq_function=acqf,
          )
          self.assertTrue(torch.equal(X, X_train))
          # test that we fail back to train_inputs if X_baseline is an empty tensor
          acqf.register_buffer("X_baseline", X_train[:0])
          X = get_X_baseline(
              acq_function=acqf,
          )
          self.assertTrue(torch.equal(X, X_train))
  
          # test acquisition function without X_baseline or model
          acqf = FixedFeatureAcquisitionFunction(
              acqf, d=2, columns=[0], values=[0]
          )
          with warnings.catch_warnings(
              record=True
          ) as w, settings.debug(True):
              X_rnd = get_X_baseline(
                  acq_function=acqf,
              )
              self.assertEqual(len(w), 1)
              self.assertTrue(
                  issubclass(w[-1].category, BotorchWarning)
              )
              self.assertIsNone(X_rnd)
  
          Y_train = 2 * X_train[:2] + 1
          moo_model = MockModel(
              MockPosterior(mean=Y_train, samples=Y_train)
          )
          ref_point = torch.zeros(2, **tkwargs)
          # test NEHVI with X_baseline
          acqf = qNoisyExpectedHypervolumeImprovement(
              moo_model,
              ref_point=ref_point,
              X_baseline=X_train[:2],
              sampler=IIDNormalSampler(sample_shape=torch.Size([2])),
              cache_root=False,
          )
          X = get_X_baseline(
              acq_function=acqf,
          )
          self.assertTrue(torch.equal(X, acqf.X_baseline))
          # test qEHVI without train_inputs
          acqf = qExpectedHypervolumeImprovement(
              moo_model,
              ref_point=ref_point,
              partitioning=FastNondominatedPartitioning(
                  ref_point=ref_point,
                  Y=Y_train,
              ),
          )
          # test extracting train_inputs from model list GP
          model_list = ModelListGP(
              SingleTaskGP(X_train, Y_train[:, :1]),
              SingleTaskGP(X_train, Y_train[:, 1:]),
          )
          acqf = qExpectedHypervolumeImprovement(
              model_list,
              ref_point=ref_point,
              partitioning=FastNondominatedPartitioning(
                  ref_point=ref_point,
                  Y=Y_train,
              ),
          )
          X = get_X_baseline(
              acq_function=acqf,
          )
          self.assertTrue(torch.equal(X, X_train))
  
          # test MESMO for which we need to use
          # `acqf.mo_model`
          batched_mo_model = SingleTaskGP(X_train, Y_train)
          acqf = qMultiObjectiveMaxValueEntropy(
              batched_mo_model,
              sample_pareto_frontiers=lambda model: torch.rand(
                  10, 2, **tkwargs
              ),
          )
          X = get_X_baseline(
              acq_function=acqf,
          )
          self.assertTrue(torch.equal(X, X_train))
          # test that if there is an input transform that is applied
          # to the train_inputs when the model is in eval mode, we
          # extract the untransformed train_inputs
          model = SingleTaskGP(
              X_train,
              Y_train[:, :1],
              input_transform=Warp(indices=[0, 1]),
          )
          model.eval()
          self.assertFalse(torch.equal(model.train_inputs[0], X_train))
          acqf = qExpectedImprovement(model, best_f=0.0)
          X = get_X_baseline(
              acq_function=acqf,
          )
          self.assertTrue(torch.equal(X, X_train))
  
imports: from __future__ import annotations
  import warnings
  import torch
  from botorch import settings
  from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction
  from botorch.acquisition.monte_carlo import qExpectedImprovement, qNoisyExpectedImprovement
  from botorch.acquisition.multi_objective.max_value_entropy_search import qMultiObjectiveMaxValueEntropy
  from botorch.acquisition.multi_objective.monte_carlo import qExpectedHypervolumeImprovement, qNoisyExpectedHypervolumeImprovement
  from botorch.exceptions import BotorchError
  from botorch.exceptions.warnings import BotorchWarning
  from botorch.models import ModelListGP, SingleTaskGP
  from botorch.models.transforms.input import Warp
  from botorch.optim.utils import columnwise_clamp, fix_features, get_X_baseline
  from botorch.sampling.normal import IIDNormalSampler
  from botorch.utils.multi_objective.box_decompositions.non_dominated import FastNondominatedPartitioning
  from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
label_desc: Use image node_id 1 for calling the TestGetXBaseline obj's underlying member method code with example usage: obj.test_get_X_baseline() without return types