code: def test_condition_on_observations__(self):
      d = 3
      for batch_shape, m, ncat, dtype, observed_noise in (
          (torch.Size(), 2, 1, torch.float, True),
          (torch.Size([2]), 1, 2, torch.double, False),
      ):
          tkwargs = {"device": self.device, "dtype": dtype}
          train_X, train_Y = _get_random_data(
              batch_shape=batch_shape, m=m, d=d, **tkwargs
          )
          cat_dims = list(range(ncat))
          train_Yvar = (
              torch.full_like(train_Y, 0.1) if observed_noise else None
          )
          model = MixedSingleTaskGP(
              train_X=train_X,
              train_Y=train_Y,
              cat_dims=cat_dims,
              train_Yvar=train_Yvar,
          )
  
          # evaluate model
          model.posterior(torch.rand(torch.Size([4, d]), **tkwargs))
          # test condition_on_observations
          fant_shape = torch.Size([2])
  
          # fantasize at different input points
          X_fant, Y_fant = _get_random_data(
              fant_shape + batch_shape, m=m, d=d, n=3, **tkwargs
          )
          additional_kwargs = (
              {"noise": torch.full_like(Y_fant, 0.1)}
              if observed_noise
              else {}
          )
          cm = model.condition_on_observations(
              X_fant, Y_fant, **additional_kwargs
          )
          # fantasize at same input points (check proper broadcasting)
          additional_kwargs = (
              {"noise": torch.full_like(Y_fant[0], 0.1)}
              if observed_noise
              else {}
          )
          cm_same_inputs = model.condition_on_observations(
              X_fant[0], Y_fant, **additional_kwargs
          )
  
          test_Xs = [
              # test broadcasting single input across fantasy and model batches
              torch.rand(4, d, **tkwargs),
              # separate input for each model batch and broadcast across
              # fantasy batches
              torch.rand(batch_shape + torch.Size([4, d]), **tkwargs),
              # separate input for each model and fantasy batch
              torch.rand(
                  fant_shape + batch_shape + torch.Size([4, d]),
                  **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.ndim > 1 else val)
                      for key, val in model.state_dict().items()
                  }
                  model_kwargs_non_batch = {
                      "train_X": train_X[0],
                      "train_Y": train_Y[0],
                      "cat_dims": cat_dims,
                  }
                  if observed_noise:
                      model_kwargs_non_batch["train_Yvar"] = train_Yvar[
                          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, d]), **tkwargs)
                  )
                  additional_kwargs = (
                      {"noise": torch.full_like(Y_fant, 0.1)}
                      if observed_noise
                      else {}
                  )
                  cm_non_batch = (
                      model_non_batch.condition_on_observations(
                          X_fant[0][0],
                          Y_fant[:, 0, :],
                          **additional_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
  import torch
  from botorch.exceptions.warnings import OptimizationWarning
  from botorch.fit import fit_gpytorch_mll
  from botorch.models.converter import batched_to_model_list
  from botorch.models.gp_regression_mixed import MixedSingleTaskGP
  from botorch.models.kernels.categorical import CategoricalKernel
  from botorch.models.transforms import Normalize
  from botorch.posteriors import GPyTorchPosterior
  from botorch.sampling import SobolQMCNormalSampler
  from botorch.utils.datasets import SupervisedDataset
  from botorch.utils.test_helpers import get_pvar_expected
  from botorch.utils.testing import _get_random_data, BotorchTestCase
  from gpytorch.kernels.kernel import AdditiveKernel, ProductKernel
  from gpytorch.kernels.matern_kernel import MaternKernel
  from gpytorch.kernels.scale_kernel import ScaleKernel
  from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
  from gpytorch.likelihoods.gaussian_likelihood import GaussianLikelihood
  from gpytorch.means import ConstantMean
  from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
label_desc: Use image node_id 2 for calling the TestMixedSingleTaskGP obj's underlying member method code with example usage: obj.test_condition_on_observations__() without return types