code: def test_KnowledgeGradient(self) -> None:
      model = KnowledgeGradient()
      model.fit(
          datasets=[self.dataset],
          metric_names=self.metric_names,
          search_space_digest=self.search_space_digest,
      )
  
      n = 2
  
      X_dummy = torch.rand(1, n, 4, **self.tkwargs)
      acq_dummy = torch.tensor(0.0, **self.tkwargs)
  
      torch_opt_config = TorchOptConfig(
          objective_weights=self.objective_weights,
          model_gen_options={
              "acquisition_function_kwargs": self.acq_options,
              "optimizer_kwargs": self.optimizer_options,
          },
      )
  
      with mock.patch(self.optimize_acqf) as mock_optimize_acqf:
          mock_optimize_acqf.side_effect = [(X_dummy, acq_dummy)]
          gen_results = model.gen(
              n=n,
              search_space_digest=self.search_space_digest,
              torch_opt_config=torch_opt_config,
          )
          self.assertTrue(
              torch.equal(gen_results.points, X_dummy.cpu())
          )
          self.assertTrue(
              torch.equal(
                  gen_results.weights,
                  torch.ones(n, dtype=self.tkwargs["dtype"]),
              )
          )
  
          # called once, the best point call is not caught by mock
          mock_optimize_acqf.assert_called_once()
  
      ini_dummy = torch.rand(10, 32, 3, **self.tkwargs)
      optimizer_options2 = {
          "num_restarts": 1,
          "raw_samples": 1,
          "options": {
              "maxiter": 5,
              "batch_limit": 1,
          },
          "partial_restarts": 2,
      }
      torch_opt_config.model_gen_options[
          "optimizer_kwargs"
      ] = optimizer_options2
      with mock.patch(
          "ax.models.torch.botorch_kg.gen_one_shot_kg_initial_conditions",
          return_value=ini_dummy,
      ) as mock_warmstart_initialization:
          gen_results = model.gen(
              n=n,
              search_space_digest=self.search_space_digest,
              torch_opt_config=torch_opt_config,
          )
          mock_warmstart_initialization.assert_called_once()
  
      posterior_tf = ScalarizedPosteriorTransform(
          weights=self.objective_weights
      )
      dummy_acq = PosteriorMean(
          model=model.model, posterior_transform=posterior_tf
      )
      with mock.patch(
          "ax.models.torch.utils.PosteriorMean", return_value=dummy_acq
      ) as mock_posterior_mean, mock.patch(
          "ax.models.torch.utils.get_botorch_objective_and_transform",
          return_value=(
              None,
              ScalarizedPosteriorTransform(
                  weights=torch.tensor([1.0], dtype=torch.double)
              ),
          ),
      ):
          gen_results = model.gen(
              n=n,
              search_space_digest=self.search_space_digest,
              torch_opt_config=torch_opt_config,
          )
          self.assertEqual(mock_posterior_mean.call_count, 2)
  
      # Check best point selection within bounds (some numerical tolerance)
      xbest = model.best_point(
          search_space_digest=self.search_space_digest,
          torch_opt_config=torch_opt_config,
      )
      lb = torch.tensor([b[0] for b in self.bounds]) - 1e-5
      ub = torch.tensor([b[1] for b in self.bounds]) + 1e-5
      self.assertTrue(torch.all(xbest <= ub))
      self.assertTrue(torch.all(xbest >= lb))
  
      # test error message
      torch_opt_config = dataclasses.replace(
          torch_opt_config,
          linear_constraints=(
              torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
              torch.tensor([[0.5], [1.0]]),
          ),
      )
      with self.assertRaises(UnsupportedError):
          gen_results = model.gen(
              n=n,
              search_space_digest=self.search_space_digest,
              torch_opt_config=torch_opt_config,
          )
  
      # test input warping
      self.assertFalse(model.use_input_warping)
      model = KnowledgeGradient(use_input_warping=True)
      model.fit(
          datasets=[self.dataset],
          search_space_digest=self.search_space_digest,
          metric_names=self.metric_names,
      )
      self.assertTrue(model.use_input_warping)
      self.assertTrue(hasattr(model.model, "input_transform"))
      self.assertIsInstance(model.model.input_transform, Warp)
  
      # test loocv pseudo likelihood
      self.assertFalse(model.use_loocv_pseudo_likelihood)
      model = KnowledgeGradient(use_loocv_pseudo_likelihood=True)
      model.fit(
          datasets=[self.dataset],
          search_space_digest=self.search_space_digest,
          metric_names=self.metric_names,
      )
      self.assertTrue(model.use_loocv_pseudo_likelihood)
  
imports: import dataclasses
  from typing import Any, Dict
  from unittest import mock
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.models.torch.botorch_kg import _instantiate_KG, KnowledgeGradient
  from ax.models.torch_base import TorchOptConfig
  from ax.utils.common.testutils import TestCase
  from ax.utils.testing.mock import fast_botorch_optimize
  from botorch.acquisition.analytic import PosteriorMean
  from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction
  from botorch.acquisition.knowledge_gradient import qMultiFidelityKnowledgeGradient
  from botorch.acquisition.monte_carlo import qSimpleRegret
  from botorch.acquisition.objective import LinearMCObjective, ScalarizedPosteriorTransform
  from botorch.exceptions.errors import UnsupportedError
  from botorch.models.transforms.input import Warp
  from botorch.sampling.normal import IIDNormalSampler, SobolQMCNormalSampler
  from botorch.utils.datasets import SupervisedDataset
label_desc: Use image node_id 2 for calling the KnowledgeGradientTest obj's underlying member method code with example usage: obj.test_KnowledgeGradient() without return types