code: def test_ALEBOGP(self) -> None:
      # First non-batch
      B = torch.tensor(
          [[1.0, 2.0, 3.0, 4.0, 5.0], [2.0, 3.0, 4.0, 5.0, 6.0]],
          dtype=torch.double,
      )
      train_X = torch.tensor(
          [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.double
      )
      train_Y = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.double)
      train_Yvar = 0.1 * torch.ones(3, 1, dtype=torch.double)
  
      with self.assertWarnsRegex(
          DeprecationWarning,
          "`get_map_model` from ax.models.torch.alebo.py is deprecated",
      ):
          mll = get_map_model(
              B=B,
              train_X=train_X,
              train_Y=train_Y,
              train_Yvar=train_Yvar,
              restarts=1,
              init_state_dict=None,
          )
      m = mll.model
      m.eval()
      self.assertIsInstance(m, ALEBOGP)
      self.assertIsInstance(m.covar_module.base_kernel, ALEBOKernel)
  
      X = torch.tensor(
          [[2.0, 2.0], [3.0, 3.0], [4.0, 4.0]], dtype=torch.double
      )
      f = m(X)
      self.assertEqual(f.mean.shape, torch.Size([3]))
      self.assertEqual(f.variance.shape, torch.Size([3]))
      self.assertEqual(f.covariance_matrix.shape, torch.Size([3, 3]))
  
      # Batch
      Uvec_b = m.covar_module.base_kernel.Uvec.repeat(5, 1)
      mean_b = m.mean_module.constant.repeat(5)
      output_scale_b = m.covar_module.raw_outputscale.repeat(5)
      m_b = get_batch_model(
          B=B,
          train_X=train_X,
          train_Y=train_Y,
          train_Yvar=train_Yvar,
          Uvec_batch=Uvec_b,
          mean_constant_batch=mean_b,
          output_scale_batch=output_scale_b,
      )
  
      self.assertEqual(m_b._aug_batch_shape, torch.Size([5]))
      f = m_b(X)
      self.assertEqual(f.mean.shape, torch.Size([3]))
      self.assertEqual(f.variance.shape, torch.Size([3]))
      self.assertEqual(f.covariance_matrix.shape, torch.Size([3, 3]))
      self.assertEqual(
          m_b.posterior(X).mvn.covariance_matrix.shape,
          torch.Size([3, 3]),
      )
  
      # The whole process in get_fitted_model
      init_state_dict = m.state_dict()
      with self.assertWarnsRegex(
          DeprecationWarning,
          "`get_fitted_model` from ax.models.torch.alebo.py is deprecated",
      ):
          m_b2 = get_fitted_model(
              B=B,
              train_X=train_X,
              train_Y=train_Y,
              train_Yvar=train_Yvar,
              restarts=1,
              nsamp=5,
              init_state_dict=init_state_dict,
          )
      self.assertEqual(m_b2._aug_batch_shape, torch.Size([5]))
  
      # Test extract_map_statedict
      map_sds = extract_map_statedict(m_b=m_b, num_outputs=1)
      self.assertEqual(len(map_sds), 1)
      self.assertEqual(len(map_sds[0]), 5)
      self.assertEqual(
          set(map_sds[0]),
          {
              "covar_module.base_kernel.Uvec",
              "covar_module.raw_outputscale",
              "mean_module.raw_constant",
              "covar_module.raw_outputscale_constraint.lower_bound",
              "covar_module.raw_outputscale_constraint.upper_bound",
          },
      )
      self.assertEqual(
          map_sds[0]["covar_module.base_kernel.Uvec"].shape,
          torch.Size([3]),
      )
  
      ml = ModelListGP(m_b, m_b2)
      map_sds = extract_map_statedict(m_b=ml, num_outputs=2)
      self.assertEqual(len(map_sds), 2)
      for i in range(2):
          self.assertEqual(len(map_sds[i]), 5)
          self.assertEqual(
              set(map_sds[i]),
              {
                  "covar_module.base_kernel.Uvec",
                  "covar_module.raw_outputscale",
                  "mean_module.raw_constant",
                  "covar_module.raw_outputscale_constraint.lower_bound",
                  "covar_module.raw_outputscale_constraint.upper_bound",
              },
          )
          self.assertEqual(
              map_sds[i]["covar_module.base_kernel.Uvec"].shape,
              torch.Size([3]),
          )
  
file_path: Ax/ax/models/tests/test_alebo.py
imports: from unittest import mock
  import numpy
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.models.torch.alebo import ALEBO, alebo_acqf_optimizer, ALEBOGP, ALEBOKernel, ei_or_nei, extract_map_statedict, get_batch_model, get_fitted_model, get_map_model
  from ax.models.torch_base import TorchOptConfig
  from ax.utils.common.testutils import TestCase
  from ax.utils.common.typeutils import checked_cast
  from ax.utils.testing.mock import fast_botorch_optimize
  from botorch.acquisition import qLogNoisyExpectedImprovement
  from botorch.acquisition.analytic import ExpectedImprovement
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.utils.datasets import SupervisedDataset
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 2 for calling the ALEBOTest obj's underlying member method code with example usage: obj.test_ALEBOGP() without return types