code: def test_condition_on_observations(self):
      for (
          iteration_fidelity,
          data_fidelities,
      ) in self.FIDELITY_TEST_PAIRS:
          n_fidelity = iteration_fidelity is not None
          if data_fidelities is not None:
              n_fidelity += len(data_fidelities)
          num_dim = 1 + n_fidelity
          for batch_shape, m, dtype, lin_trunc in itertools.product(
              (torch.Size(), torch.Size([2])),
              (1, 2),
              (torch.float, torch.double),
              (False, True),
          ):
              tkwargs = {"device": self.device, "dtype": dtype}
              model, model_kwargs = self._get_model_and_data(
                  iteration_fidelity=iteration_fidelity,
                  data_fidelities=data_fidelities,
                  batch_shape=batch_shape,
                  m=m,
                  lin_truncated=lin_trunc,
                  **tkwargs,
              )
              # evaluate model
              model.posterior(
                  torch.rand(torch.Size([4, num_dim]), **tkwargs)
              )
              # test condition_on_observations
              fant_shape = torch.Size([2])
              # fantasize at different input points
              X_fant, Y_fant = _get_random_data_with_fidelity(
                  fant_shape + batch_shape,
                  m,
                  n_fidelity=n_fidelity,
                  n=3,
                  **tkwargs,
              )
              c_kwargs = (
                  {"noise": torch.full_like(Y_fant, 0.01)}
                  if isinstance(
                      model.likelihood, FixedNoiseGaussianLikelihood
                  )
                  else {}
              )
              cm = model.condition_on_observations(
                  X_fant, Y_fant, **c_kwargs
              )
              # fantasize at different same input points
              c_kwargs_same_inputs = (
                  {"noise": torch.full_like(Y_fant[0], 0.01)}
                  if isinstance(
                      model.likelihood, FixedNoiseGaussianLikelihood
                  )
                  else {}
              )
              cm_same_inputs = model.condition_on_observations(
                  X_fant[0], Y_fant, **c_kwargs_same_inputs
              )
  
              test_Xs = [
                  # test broadcasting single input across fantasy and
                  # model batches
                  torch.rand(4, num_dim, **tkwargs),
                  # separate input for each model batch and broadcast across
                  # fantasy batches
                  torch.rand(
                      batch_shape + torch.Size([4, num_dim]), **tkwargs
                  ),
                  # separate input for each model and fantasy batch
                  torch.rand(
                      fant_shape
                      + batch_shape
                      + torch.Size([4, num_dim]),
                      **tkwargs,
                  ),
              ]
              for test_X in test_Xs:
                  posterior = cm.posterior(test_X)
                  self.assertEqual(
                      posterior.mean.shape,
                      fant_shape + batch_shape + torch.Size([4, m]),
                  )
                  posterior_same_inputs = cm_same_inputs.posterior(
                      test_X
                  )
                  self.assertEqual(
                      posterior_same_inputs.mean.shape,
                      fant_shape + batch_shape + torch.Size([4, m]),
                  )
  
                  # check that fantasies of batched model are correct
                  if len(batch_shape) > 0 and test_X.dim() == 2:
                      state_dict_non_batch = {
                          key: (val[0] if val.numel() > 1 else val)
                          for key, val in model.state_dict().items()
                      }
  
                      model_kwargs_non_batch = {}
                      for k, v in model_kwargs.items():
                          if k in (
                              "iteration_fidelity",
                              "data_fidelities",
                              "linear_truncated",
                              "input_transform",
                          ):
                              model_kwargs_non_batch[k] = v
                          else:
                              model_kwargs_non_batch[k] = v[0]
  
                      model_non_batch = type(model)(
                          **model_kwargs_non_batch
                      )
                      model_non_batch.load_state_dict(
                          state_dict_non_batch
                      )
                      model_non_batch.eval()
                      model_non_batch.likelihood.eval()
                      model_non_batch.posterior(
                          torch.rand(
                              torch.Size([4, num_dim]), **tkwargs
                          )
                      )
                      c_kwargs = (
                          {
                              "noise": torch.full_like(
                                  Y_fant[0, 0, :], 0.01
                              )
                          }
                          if isinstance(
                              model.likelihood,
                              FixedNoiseGaussianLikelihood,
                          )
                          else {}
                      )
                      mnb = model_non_batch
                      cm_non_batch = mnb.condition_on_observations(
                          X_fant[0][0], Y_fant[:, 0, :], **c_kwargs
                      )
                      non_batch_posterior = cm_non_batch.posterior(
                          test_X
                      )
                      self.assertTrue(
                          torch.allclose(
                              posterior_same_inputs.mean[:, 0, ...],
                              non_batch_posterior.mean,
                              atol=1e-3,
                          )
                      )
                      self.assertTrue(
                          torch.allclose(
                              posterior_same_inputs.distribution.covariance_matrix[
                                  :, 0, :, :
                              ],
                              non_batch_posterior.distribution.covariance_matrix,
                              atol=1e-3,
                          )
                      )
  
imports: import itertools
  import warnings
  from typing import Tuple
  import torch
  from botorch.exceptions.errors import UnsupportedError
  from botorch.exceptions.warnings import OptimizationWarning
  from botorch.fit import fit_gpytorch_mll
  from botorch.models.gp_regression_fidelity import FixedNoiseMultiFidelityGP, SingleTaskMultiFidelityGP
  from botorch.models.transforms import Normalize, Standardize
  from botorch.posteriors import GPyTorchPosterior
  from botorch.sampling import SobolQMCNormalSampler
  from botorch.utils.datasets import SupervisedDataset
  from botorch.utils.testing import _get_random_data, BotorchTestCase
  from gpytorch.kernels.scale_kernel import ScaleKernel
  from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
  from gpytorch.means import ConstantMean
  from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
  from torch import Tensor
label_desc: Use image node_id 4 for calling the TestSingleTaskMultiFidelityGP obj's underlying member method code with example usage: obj.test_condition_on_observations() without return types