code: def test_cache_root(self):
      sample_cached_path = (
          "botorch.acquisition.cached_cholesky.sample_cached_cholesky"
      )
      raw_state_dict = {
          "likelihood.noise_covar.raw_noise": torch.tensor(
              [[0.0895], [0.2594]], dtype=torch.float64
          ),
          "mean_module.raw_constant": torch.tensor(
              [-0.4545, -0.1285], dtype=torch.float64
          ),
          "covar_module.raw_outputscale": torch.tensor(
              [1.4876, 1.4897], dtype=torch.float64
          ),
          "covar_module.base_kernel.raw_lengthscale": torch.tensor(
              [[[-0.7202, -0.2868]], [[-0.8794, -1.2877]]],
              dtype=torch.float64,
          ),
      }
      # test batched models (e.g. for MCMC)
      for train_batch_shape, m, dtype in product(
          (torch.Size([]), torch.Size([3])),
          (1, 2),
          (torch.float, torch.double),
      ):
          state_dict = deepcopy(raw_state_dict)
          for k, v in state_dict.items():
              if m == 1:
                  v = v[0]
              if len(train_batch_shape) > 0:
                  v = v.unsqueeze(0).expand(
                      *train_batch_shape, *v.shape
                  )
              state_dict[k] = v
          tkwargs = {"device": self.device, "dtype": dtype}
          if m == 2:
              objective = GenericMCObjective(lambda Y, X: Y.sum(dim=-1))
          else:
              objective = None
          for k, v in state_dict.items():
              state_dict[k] = v.to(**tkwargs)
          all_close_kwargs = (
              {
                  "atol": 1e-1,
                  "rtol": 0.0,
              }
              if dtype == torch.float
              else {"atol": 1e-4, "rtol": 0.0}
          )
          torch.manual_seed(1234)
          train_X = torch.rand(*train_batch_shape, 3, 2, **tkwargs)
          train_Y = (
              torch.sin(train_X * 2 * pi)
              + torch.randn(*train_batch_shape, 3, 2, **tkwargs)
          )[..., :m]
          train_Y = standardize(train_Y)
          model = SingleTaskGP(
              train_X,
              train_Y,
          )
          if len(train_batch_shape) > 0:
              X_baseline = train_X[0]
          else:
              X_baseline = train_X
          model.load_state_dict(state_dict, strict=False)
          sampler = IIDNormalSampler(
              sample_shape=torch.Size([5]), seed=0
          )
          torch.manual_seed(0)
          acqf = qNoisyExpectedImprovement(
              model=model,
              X_baseline=X_baseline,
              sampler=sampler,
              objective=objective,
              prune_baseline=False,
              cache_root=True,
          )
  
          orig_base_samples = (
              acqf.base_sampler.base_samples.detach().clone()
          )
          sampler2 = IIDNormalSampler(
              sample_shape=torch.Size([5]), seed=0
          )
          sampler2.base_samples = orig_base_samples
          torch.manual_seed(0)
          acqf_no_cache = qNoisyExpectedImprovement(
              model=model,
              X_baseline=X_baseline,
              sampler=sampler2,
              objective=objective,
              prune_baseline=False,
              cache_root=False,
          )
          for q, batch_shape in product(
              (1, 3),
              (torch.Size([]), torch.Size([3]), torch.Size([4, 3])),
          ):
              acqf.q_in = -1
              acqf_no_cache.q_in = -1
              test_X = (
                  0.3
                  + 0.05 * torch.randn(*batch_shape, q, 2, **tkwargs)
              ).requires_grad_(True)
              with mock.patch(
                  sample_cached_path, wraps=sample_cached_cholesky
              ) as mock_sample_cached:
                  torch.manual_seed(0)
                  val = acqf(test_X)
                  mock_sample_cached.assert_called_once()
              val.sum().backward()
              base_samples = acqf.sampler.base_samples.detach().clone()
              X_grad = test_X.grad.clone()
              test_X2 = test_X.detach().clone().requires_grad_(True)
              acqf_no_cache.sampler.base_samples = base_samples
              with mock.patch(
                  sample_cached_path, wraps=sample_cached_cholesky
              ) as mock_sample_cached:
                  torch.manual_seed(0)
                  val2 = acqf_no_cache(test_X2)
                  mock_sample_cached.assert_not_called()
              self.assertAllClose(val, val2, **all_close_kwargs)
              val2.sum().backward()
              self.assertAllClose(
                  X_grad, test_X2.grad, **all_close_kwargs
              )
          # test we fall back to standard sampling for
          # ill-conditioned covariances
          acqf._baseline_L = torch.zeros_like(acqf._baseline_L)
          with warnings.catch_warnings(
              record=True
          ) as ws, settings.debug(True):
              with torch.no_grad():
                  acqf(test_X)
          self.assertEqual(
              sum(issubclass(w.category, BotorchWarning) for w in ws), 1
          )
  
      # test w/ posterior transform
      X_baseline = torch.rand(2, 1)
      model = SingleTaskGP(X_baseline, torch.randn(2, 1))
      pt = ScalarizedPosteriorTransform(weights=torch.tensor([-1]))
      with mock.patch.object(
          qNoisyExpectedImprovement,
          "_compute_root_decomposition",
      ) as mock_cache_root:
          acqf = qNoisyExpectedImprovement(
              model=model,
              X_baseline=X_baseline,
              sampler=IIDNormalSampler(sample_shape=torch.Size([1])),
              posterior_transform=pt,
              prune_baseline=False,
              cache_root=True,
          )
          tf_post = model.posterior(X_baseline, posterior_transform=pt)
          self.assertTrue(
              torch.allclose(
                  tf_post.mean,
                  mock_cache_root.call_args[-1]["posterior"].mean,
              )
          )
  
      # testing constraints
      n, d, m = 8, 1, 3
      X_baseline = torch.rand(n, d)
      model = SingleTaskGP(
          X_baseline, torch.randn(n, m)
      )  # batched model
      nei_args = {
          "model": model,
          "X_baseline": X_baseline,
          "prune_baseline": False,
          "cache_root": True,
          "posterior_transform": ScalarizedPosteriorTransform(
              weights=torch.ones(m)
          ),
          "sampler": SobolQMCNormalSampler(
              sample_shape=torch.Size([5])
          ),
      }
      acqf = qNoisyExpectedImprovement(**nei_args)
      X = torch.randn_like(X_baseline)
      for con in [feasible_con, infeasible_con]:
          with self.subTest(con=con):
              target = "botorch.acquisition.utils.get_infeasible_cost"
              infcost = torch.tensor(
                  [3], device=self.device, dtype=dtype
              )
              with mock.patch(target, return_value=infcost):
                  cacqf = qNoisyExpectedImprovement(
                      **nei_args, constraints=[con]
                  )
  
              _, obj = cacqf._get_samples_and_objectives(X)
              best_feas_f = cacqf.compute_best_f(obj)
              if con is feasible_con:
                  self.assertAllClose(
                      best_feas_f, acqf.compute_best_f(obj)
                  )
              else:
                  self.assertAllClose(
                      best_feas_f,
                      torch.full_like(obj[..., [0]], -infcost.item()),
                  )
  
imports: import warnings
  from copy import deepcopy
  from functools import partial
  from itertools import product
  from math import pi
  from unittest import mock
  import torch
  from botorch import settings
  from botorch.acquisition.monte_carlo import MCAcquisitionFunction, qExpectedImprovement, qNoisyExpectedImprovement, qProbabilityOfImprovement, qSimpleRegret, qUpperConfidenceBound, SampleReducingMCAcquisitionFunction
  from botorch.acquisition.objective import ConstrainedMCObjective, GenericMCObjective, IdentityMCObjective, ScalarizedPosteriorTransform
  from botorch.acquisition.utils import prune_inferior_points
  from botorch.exceptions import BotorchWarning, UnsupportedError
  from botorch.models import SingleTaskGP
  from botorch.sampling.normal import IIDNormalSampler, SobolQMCNormalSampler
  from botorch.utils.low_rank import sample_cached_cholesky
  from botorch.utils.test_helpers import DummyNonScalarizingPosteriorTransform
  from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
  from botorch.utils.transforms import standardize
  from torch import Tensor
label_desc: Use image node_id 4 for calling the TestQNoisyExpectedImprovement obj's underlying member method code with example usage: obj.test_cache_root() without return types