code: def test_FullyBayesianBotorchModel(
      self, dtype: torch.dtype = torch.float, cuda: bool = False
  ) -> None:
      # test deprecation warning
      warnings.resetwarnings()  # this is necessary for building in mode/opt
      warnings.simplefilter("always", append=True)
      with warnings.catch_warnings(record=True) as ws:
          self.model_cls(use_saas=True)
          # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
          #  attribute `assertTrue`.
          self.assertTrue(
              any(
                  issubclass(w.category, DeprecationWarning) for w in ws
              )
          )
          msg = "Passing `use_saas` is no longer supported"
          self.assertTrue(any(msg in str(w.message) for w in ws))
      Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = get_torch_test_data(
          dtype=dtype, cuda=cuda, constant_noise=True
      )
      Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(
          dtype=dtype, cuda=cuda, constant_noise=True
      )
      Yvars_inferred_noise = [
          torch.full_like(Yvars1[0], float("nan")),
          torch.full_like(Yvars2[0], float("nan")),
      ]
      # make input different for each output
      Xs2_diff = [Xs2[0] + 0.1]
      Xs = Xs1 + Xs2_diff
      Ys = Ys1 + Ys2
  
      options = product([True, False], [True, False], ["matern", "rbf"])
      for inferred_noise, use_input_warping, gp_kernel in options:
          Yvars = (
              Yvars_inferred_noise
              if inferred_noise
              else Yvars1 + Yvars2
          )
          model = self.model_cls(
              use_input_warping=use_input_warping,
              thinning=1,
              num_samples=4,
              disable_progbar=True,
              max_tree_depth=1,
              gp_kernel=gp_kernel,
              verbose=True,
          )
          if use_input_warping:
              self.assertTrue(model.use_input_warping)
          # Test ModelListGP
          # make training data different for each output
          tkwargs: Dict[str, Any] = {
              "dtype": dtype,
              "device": Xs1[0].device,
          }
          dummy_samples_list = _get_dummy_mcmc_samples(
              num_samples=4, num_outputs=2, **tkwargs
          )
          for dummy_samples in dummy_samples_list:
              if use_input_warping:
                  # pyre-fixme[16]: `str` has no attribute `__setitem__`.
                  dummy_samples["c0"] = (
                      torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs)
                      * 0.5
                      + 0.1
                  )
                  dummy_samples["c1"] = (
                      torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs)
                      * 0.5
                      + 0.1
                  )
              if inferred_noise:
                  dummy_samples["noise"] = torch.rand(
                      4, 1, **tkwargs
                  ).clamp_min(MIN_INFERRED_NOISE_LEVEL)
  
          with mock.patch(
              RUN_INFERENCE_PATH,
              side_effect=dummy_samples_list,
          ) as _mock_fit_model:
              model.fit(
                  datasets=[
                      SupervisedDataset(
                          X=X,
                          Y=Y,
                          Yvar=Yvar,
                          feature_names=fns,
                          outcome_names=[f"y{i}"],
                      )
                      for X, Y, Yvar, i in zip(Xs, Ys, Yvars, count())
                  ],
                  search_space_digest=SearchSpaceDigest(
                      feature_names=fns,
                      bounds=bounds,
                      task_features=tfs,
                  ),
                  metric_names=["y1", "y2"],
              )
  
              # Check that there are no unexpected constraints on the hyperparameters
              for _, m in enumerate(
                  cast(ModelList, model.model).models
              ):
                  if inferred_noise:
                      noise_covar = m.likelihood.noise_covar
                      self.assertEqual(  # pyre-fixme[16]
                          noise_covar.raw_noise_constraint.__class__,
                          GreaterThan,
                      )
                      self.assertTrue(
                          torch.allclose(
                              noise_covar.raw_noise_constraint.lower_bound,
                              torch.tensor(1e-4, dtype=dtype),
                          )
                      )
                  else:
                      self.assertFalse(  # pyre-fixme[16]
                          hasattr(
                              m.likelihood.noise_covar,
                              "raw_noise_constraint",
                          )
                      )
  
                  self.assertEqual(
                      m.covar_module.__class__, ScaleKernel
                  )
                  self.assertEqual(
                      m.covar_module.raw_outputscale_constraint.__class__,
                      Positive,
                  )
                  self.assertEqual(
                      m.covar_module.base_kernel.__class__,
                      RBFKernel if gp_kernel == "rbf" else MaternKernel,
                  )
                  self.assertEqual(
                      m.covar_module.base_kernel.ard_num_dims, 3
                  )
                  self.assertEqual(
                      m.covar_module.base_kernel.raw_lengthscale_constraint.__class__,
                      Positive,
                  )
  
              #  attribute `assertEqual`.
              self.assertEqual(_mock_fit_model.call_count, 2)
              for i, call in enumerate(_mock_fit_model.call_args_list):
                  _, ckwargs = call
                  X = Xs[i]
                  Y = Ys[i]
                  Yvar = Yvars[i]
                  # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
                  #  attribute `assertIs`.
                  self.assertIs(
                      ckwargs["pyro_model"], single_task_pyro_model
                  )
  
                  self.assertTrue(torch.equal(ckwargs["X"], X))
                  self.assertTrue(torch.equal(ckwargs["Y"], Y))
                  if inferred_noise:
                      self.assertTrue(
                          torch.isnan(ckwargs["Yvar"]).all()
                      )
                  else:
                      self.assertTrue(
                          torch.equal(ckwargs["Yvar"], Yvar)
                      )
                  self.assertEqual(ckwargs["num_samples"], 4)
                  self.assertEqual(ckwargs["warmup_steps"], 512)
                  self.assertEqual(ckwargs["max_tree_depth"], 1)
                  self.assertTrue(ckwargs["disable_progbar"])
                  self.assertFalse(ckwargs["jit_compile"])
                  self.assertEqual(
                      ckwargs["use_input_warping"], use_input_warping
                  )
                  self.assertEqual(ckwargs["gp_kernel"], gp_kernel)
                  self.assertTrue(ckwargs["verbose"])
  
                  # Check attributes
                  self.assertTrue(torch.equal(model.Xs[i], Xs[i]))
                  self.assertEqual(model.dtype, Xs[i].dtype)
                  self.assertEqual(model.device, Xs[i].device)
                  # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
                  #  attribute `assertIsInstance`.
                  self.assertIsInstance(model.model, ModelListGP)
  
                  # Check fitting
                  # Note each model in the model list is a batched model, where
                  # the batch dim corresponds to the MCMC samples
                  model_list = cast(ModelList, model.model).models
                  # Put model in `eval` mode to transform the train inputs.
                  m = model_list[i].eval()
                  # check mcmc samples
                  # pyre-fixme[6]: For 1st param expected `str` but got `int`.
                  dummy_samples = dummy_samples_list[i]
                  expected_train_inputs = Xs[i].expand(4, *Xs[i].shape)
                  if use_input_warping:
                      # train inputs should be warped inputs
                      expected_train_inputs = m.input_transform(
                          expected_train_inputs
                      )
                  self.assertTrue(
                      torch.equal(
                          m.train_inputs[0],
                          expected_train_inputs,
                      )
                  )
                  self.assertTrue(
                      torch.equal(
                          m.train_targets,
                          Ys[i].view(1, -1).expand(4, Ys[i].numel()),
                      )
                  )
                  expected_noise = (
                      # pyre-fixme[6]: For 1st param expected `Union[None,
                      #  List[typing.Any], int, slice, Tensor,
                      #  typing.Tuple[typing.Any, ...]]` but got `str`.
                      dummy_samples["noise"].view(
                          m.likelihood.noise.shape
                      )
                      if inferred_noise
                      else Yvars[i]
                      .view(1, -1)
                      .expand(4, Yvars[i].numel())
                  )
                  self.assertTrue(
                      torch.allclose(
                          m.likelihood.noise.detach(),
                          expected_noise,
                      )
                  )
                  self.assertIsInstance(
                      m.likelihood, _GaussianLikelihoodBase
                  )
                  self.assertTrue(
                      torch.allclose(
                          m.covar_module.base_kernel.lengthscale.detach(),
                          # pyre-fixme[6]: For 1st param expected `Union[None,
                          #  List[typing.Any], int, slice, Tensor,
                          #  typing.Tuple[typing.Any, ...]]` but got `str`.
                          dummy_samples["lengthscale"].view(
                              m.covar_module.base_kernel.lengthscale.shape
                          ),
                          rtol=1e-4,
                          atol=1e-6,
                      )
                  )
                  self.assertTrue(
                      torch.allclose(
                          m.covar_module.outputscale.detach(),
                          # pyre-fixme[6]: For 1st param expected `Union[None,
                          #  List[typing.Any], int, slice, Tensor,
                          #  typing.Tuple[typing.Any, ...]]` but got `str`.
                          dummy_samples["outputscale"].view(
                              m.covar_module.outputscale.shape
                          ),
                      )
                  )
                  self.assertTrue(
                      torch.allclose(
                          m.mean_module.constant.detach(),
                          # pyre-fixme[6]: For 1st param expected `Union[None,
                          #  List[typing.Any], int, slice, Tensor,
                          #  typing.Tuple[typing.Any, ...]]` but got `str`.
                          dummy_samples["mean"].view(
                              m.mean_module.constant.shape
                          ),
                      )
                  )
                  if use_input_warping:
                      self.assertTrue(hasattr(m, "input_transform"))
                      self.assertIsInstance(m.input_transform, Warp)
                      self.assertTrue(
                          torch.equal(
                              m.input_transform.concentration0,
                              # pyre-fixme[6]: For 1st param expected `str`
                              #  but got `int`.
                              # pyre-fixme[6]: For 1st param expected
                              #  `Union[None, List[typing.Any], int, slice,
                              #  Tensor, typing.Tuple[typing.Any, ...]]` but got
                              #  `str`.
                              dummy_samples_list[i]["c0"],
                          )
                      )
                      self.assertTrue(
                          torch.equal(
                              m.input_transform.concentration1,
                              # pyre-fixme[6]: For 1st param expected `str`
                              #  but got `int`.
                              # pyre-fixme[6]: For 1st param expected
                              #  `Union[None, List[typing.Any], int, slice,
                              #  Tensor, typing.Tuple[typing.Any, ...]]` but got
                              #  `str`.
                              dummy_samples_list[i]["c1"],
                          )
                      )
                  else:
                      self.assertFalse(hasattr(m, "input_transform"))
          # test that multi-task is not implemented
          (
              Xs_mt,
              Ys_mt,
              Yvars_mt,
              bounds_mt,
              tfs_mt,
              fns_mt,
              mns_mt,
          ) = get_torch_test_data(
              dtype=dtype,
              cuda=cuda,
              constant_noise=True,
              task_features=[2],
          )
          with mock.patch(
              RUN_INFERENCE_PATH,
              side_effect=dummy_samples_list,
              # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
              #  attribute `assertRaises`.
          ) as _mock_fit_model, self.assertRaises(NotImplementedError):
              model.fit(
                  datasets=[
                      SupervisedDataset(
                          X=X,
                          Y=Y,
                          Yvar=Yvar,
                          feature_names=fns_mt,
                          outcome_names=[mn],
                      )
                      for X, Y, Yvar, mn in zip(
                          Xs_mt, Ys_mt, Yvars_mt, mns_mt
                      )
                  ],
                  metric_names=mns_mt,
                  search_space_digest=SearchSpaceDigest(
                      feature_names=fns_mt,
                      bounds=bounds_mt,
                      task_features=tfs_mt,
                  ),
              )
          with mock.patch(
              RUN_INFERENCE_PATH,
              side_effect=dummy_samples_list,
          ) as _mock_fit_model, self.assertRaises(NotImplementedError):
              model.fit(
                  datasets=[
                      SupervisedDataset(
                          X=X,
                          Y=Y,
                          Yvar=Yvar,
                          feature_names=fns,
                          outcome_names=[mn],
                      )
                      for X, Y, Yvar, mn in zip(
                          Xs1 + Xs2, Ys1 + Ys2, Yvars1 + Yvars2, mns * 2
                      )
                  ],
                  metric_names=mns,
                  search_space_digest=SearchSpaceDigest(
                      feature_names=fns,
                      bounds=bounds,
                      fidelity_features=[0],
                  ),
              )
          # fit model with same inputs (otherwise X_observed will be None)
          model = self.model_cls(
              use_input_warping=use_input_warping,
              thinning=1,
              num_samples=4,
              disable_progbar=True,
              max_tree_depth=1,
              gp_kernel=gp_kernel,
          )
          Yvars = Yvars1 + Yvars2
          dummy_samples_list = _get_dummy_mcmc_samples(
              num_samples=4, num_outputs=2, **tkwargs
          )
          with mock.patch(
              RUN_INFERENCE_PATH,
              side_effect=dummy_samples_list,
          ) as _mock_fit_model:
              model.fit(
                  datasets=[
                      SupervisedDataset(
                          X=X,
                          Y=Y,
                          Yvar=Yvar,
                          feature_names=fns,
                          outcome_names=[mn],
                      )
                      for X, Y, Yvar, mn in zip(
                          Xs1 + Xs2, Ys1 + Ys2, Yvars, mns * 2
                      )
                  ],
                  metric_names=mns,
                  search_space_digest=SearchSpaceDigest(
                      feature_names=fns,
                      bounds=bounds,
                      task_features=tfs,
                  ),
              )
  
          # Check the hyperparameters and shapes
          models = cast(ModelList, model.model).models
          self.assertEqual(len(models), 2)
          m1, m2 = models
          # Mean
          self.assertEqual(m1.mean_module.constant.shape, (4,))
          self.assertFalse(
              torch.isclose(
                  m1.mean_module.constant, m2.mean_module.constant
              ).any()
          )
          # Outputscales
          self.assertEqual(m1.covar_module.outputscale.shape, (4,))
          self.assertFalse(
              torch.isclose(
                  m1.covar_module.outputscale,
                  m2.covar_module.outputscale,
              ).any()
          )
          # Lengthscales
          self.assertEqual(
              m1.covar_module.base_kernel.lengthscale.shape, (4, 1, 3)
          )
          self.assertFalse(
              torch.isclose(
                  m1.covar_module.base_kernel.lengthscale,
                  m2.covar_module.base_kernel.lengthscale,
              ).any()
          )
  
          # Check infeasible cost can be computed on the model
          device = torch.device("cuda") if cuda else torch.device("cpu")
          objective_weights = torch.tensor(
              [1.0, 0.0], dtype=dtype, device=device
          )
          objective_transform = get_objective_weights_transform(
              objective_weights
          )
          infeasible_cost = (
              get_infeasible_cost(
                  X=Xs1[0],
                  model=model.model,
                  objective=objective_transform,
              )
              .detach()
              .clone()
          )
          posterior = cast(
              GPyTorchPosterior, model.model.posterior(Xs1[0])
          )
          expected_infeasible_cost = -1 * torch.min(
              # pyre-fixme[20]: Argument `1` expected.
              objective_transform(
                  posterior.mean - 6 * posterior.variance.sqrt()
              ).min(),
              torch.tensor(0.0, dtype=dtype, device=device),
          )
          self.assertTrue(
              torch.abs(infeasible_cost - expected_infeasible_cost)
              < 1e-5
          )
  
          # Check prediction
          X = torch.tensor([[6.0, 7.0, 8.0]], **tkwargs)
          f_mean, f_cov = model.predict(X)
          self.assertTrue(f_mean.shape == torch.Size([1, 2]))
          self.assertTrue(f_cov.shape == torch.Size([1, 2, 2]))
  
          # Check generation
          objective_weights = torch.tensor(
              [1.0, 0.0]
              if self.model_cls is FullyBayesianBotorchModel
              else [1.0, 1.0],
              **tkwargs,
          )
          outcome_constraints = (
              torch.tensor([[0.0, 1.0]], **tkwargs),
              torch.tensor([[5.0]], **tkwargs),
          )
          linear_constraints = (
              torch.tensor([[0.0, 1.0, 1.0]]),
              torch.tensor([[100.0]]),
          )
          fixed_features = None
          pending_observations = [
              torch.tensor([[1.0, 3.0, 4.0]], **tkwargs),
              torch.tensor([[2.0, 6.0, 8.0]], **tkwargs),
          ]
          n = 3
  
          X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
          acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
          model_gen_options = {
              Keys.OPTIMIZER_KWARGS: {"options": {"maxiter": 1}},
              Keys.ACQF_KWARGS: {"mc_samples": 3},
          }
          search_space_digest = SearchSpaceDigest(
              feature_names=fns,
              bounds=bounds,
          )
          torch_opt_config = TorchOptConfig(
              objective_weights=objective_weights,
              objective_thresholds=torch.zeros(2, **tkwargs)
              if self.model_cls is FullyBayesianMOOBotorchModel
              else None,
              outcome_constraints=outcome_constraints,
              linear_constraints=linear_constraints,
              fixed_features=fixed_features,
              pending_observations=pending_observations,
              # pyre-fixme[6]: For 7th param expected `Dict[str, Union[None,
              #  Dict[str, typing.Any], OptimizationConfig, AcquisitionFunction,
              #  float, int, str]]` but got `Dict[Keys, Dict[str, int]]`.
              model_gen_options=model_gen_options,
              rounding_func=dummy_func,
              is_moo=self.model_cls is FullyBayesianMOOBotorchModel,
          )
          # test sequential optimize with constraints
          with mock.patch(
              "ax.models.torch.botorch_defaults.optimize_acqf",
              return_value=(X_dummy, acqfv_dummy),
          ) as _:
              # Xgen, wgen, gen_metadata, cand_metadata = model.gen(
              gen_results = model.gen(
                  n=n,
                  search_space_digest=search_space_digest,
                  torch_opt_config=torch_opt_config,
              )
              # note: gen() always returns CPU tensors
              self.assertTrue(
                  torch.equal(gen_results.points, X_dummy.cpu())
              )
              self.assertTrue(
                  torch.equal(
                      gen_results.weights, torch.ones(n, dtype=dtype)
                  )
              )
  
          # actually test optimization for 1 step without constraints
          with mock.patch(
              "ax.models.torch.botorch_defaults.optimize_acqf",
              wraps=optimize_acqf,
              return_value=(X_dummy, acqfv_dummy),
          ) as _:
              gen_results = model.gen(
                  n=n,
                  search_space_digest=search_space_digest,
                  torch_opt_config=dataclasses.replace(
                      torch_opt_config, linear_constraints=None
                  ),
              )
              # note: gen() always returns CPU tensors
              self.assertTrue(
                  torch.equal(gen_results.points, X_dummy.cpu())
              )
              self.assertTrue(
                  torch.equal(
                      gen_results.weights, torch.ones(n, dtype=dtype)
                  )
              )
  
          # Check best point selection
          if self.model_cls is FullyBayesianMOOBotorchModel:
              # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
              #  attribute `assertRaisesRegex`.
              with self.assertRaisesRegex(
                  NotImplementedError, "Best observed"
              ):
                  model.best_point(
                      search_space_digest=search_space_digest,
                      torch_opt_config=torch_opt_config,
                  )
          else:
              # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
              #  attribute `assertIsNotNone`.
              self.assertIsNotNone(
                  model.best_point(
                      search_space_digest=search_space_digest,
                      torch_opt_config=torch_opt_config,
                  )
              )
              # pyre-fixme[16]: `BaseFullyBayesianBotorchModelTest` has no
              #  attribute `assertIsNone`.
              self.assertIsNone(
                  model.best_point(
                      search_space_digest=search_space_digest,
                      torch_opt_config=dataclasses.replace(
                          torch_opt_config, fixed_features={0: 100.0}
                      ),
                  )
              )
  
          # Test cross-validation
          mean, variance = model.cross_validate(
              datasets=[
                  SupervisedDataset(
                      X=X,
                      Y=Y,
                      Yvar=Yvar,
                      feature_names=fns,
                      outcome_names=[mn],
                  )
                  for X, Y, Yvar, mn in zip(
                      Xs1 + Xs2, Ys1 + Ys2, Yvars, mns * 2
                  )
              ],
              X_test=torch.tensor(
                  [[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
                  dtype=dtype,
                  device=device,
              ),
          )
          self.assertTrue(mean.shape == torch.Size([2, 2]))
          self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
  
          # Test cross-validation with refit_on_cv
          model.refit_on_cv = True
          with mock.patch(
              RUN_INFERENCE_PATH,
              side_effect=dummy_samples_list,
          ) as _mock_fit_model:
              mean, variance = model.cross_validate(
                  datasets=[
                      SupervisedDataset(
                          X=X,
                          Y=Y,
                          Yvar=Yvar,
                          feature_names=fns,
                          outcome_names=[mn],
                      )
                      for X, Y, Yvar, mn in zip(
                          Xs1 + Xs2, Ys1 + Ys2, Yvars, mns * 2
                      )
                  ],
                  X_test=torch.tensor(
                      [[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
                      dtype=dtype,
                      device=device,
                  ),
              )
              self.assertTrue(mean.shape == torch.Size([2, 2]))
              self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
  
          # Test feature_importances
          importances = model.feature_importances()
          self.assertEqual(importances.shape, torch.Size([2, 1, 3]))
  
          # test unfit model CV and feature_importances
          unfit_model = self.model_cls()
          with self.assertRaises(RuntimeError):
              unfit_model.cross_validate(
                  datasets=[
                      SupervisedDataset(
                          X=X,
                          Y=Y,
                          Yvar=Yvar,
                          feature_names=fns,
                          outcome_names=[mn],
                      )
                      for X, Y, Yvar, mn in zip(
                          Xs1 + Xs2, Ys1 + Ys2, Yvars1 + Yvars2, mns * 2
                      )
                  ],
                  X_test=Xs1[0],
              )
          with self.assertRaises(RuntimeError):
              unfit_model.feature_importances()
  
imports: import dataclasses
  import warnings
  from abc import ABC
  from contextlib import ExitStack
  from itertools import count, product
  from logging import Logger
  from math import sqrt
  from typing import Any, cast, Dict, Type
  from unittest import mock
  import pyro
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.exceptions.core import AxError
  from ax.models.torch.fully_bayesian import FullyBayesianBotorchModel, FullyBayesianMOOBotorchModel, matern_kernel, rbf_kernel, single_task_pyro_model
  from ax.models.torch_base import TorchOptConfig
  from ax.utils.common.constants import Keys
  from ax.utils.common.logger import get_logger
  from ax.utils.common.testutils import TestCase
  from ax.utils.testing.torch_stubs import get_torch_test_data
  from botorch.acquisition.utils import get_infeasible_cost
  from botorch.models import ModelListGP
  from botorch.models.gp_regression import MIN_INFERRED_NOISE_LEVEL
  from botorch.models.model import ModelList
  from botorch.models.transforms.input import Warp
  from botorch.optim.optimize import optimize_acqf
  from botorch.posteriors.gpytorch import GPyTorchPosterior
  from botorch.utils import get_objective_weights_transform
  from botorch.utils.datasets import SupervisedDataset
  from gpytorch.constraints import GreaterThan, Positive
  from gpytorch.kernels import MaternKernel, RBFKernel, ScaleKernel
  from gpytorch.likelihoods import _GaussianLikelihoodBase
  from pyro.infer.mcmc import MCMC, NUTS
  from pyro.ops.integrator import potential_grad
label_desc: Use image node_id 1 for calling the BaseFullyBayesianBotorchModelTest obj's underlying member method code with example usage: obj.test_FullyBayesianBotorchModel(dtype, cuda) without return types