class_docstr: Mixin to add a `fantasize` method to a `Model`.
  
  Example:
      class BaseModel:
          def __init__(self, ...):
          def condition_on_observations(self, ...):
          def posterior(self, ...):
          def transform_inputs(self, ...):
  
      class ModelThatCanFantasize(BaseModel, FantasizeMixin):
          def __init__(self, args):
              super().__init__(args)
  
      model = ModelThatCanFantasize(...)
      model.fantasize(X)
class_docstr_tok: ["Mixin","to","add","a","`","fantasize","`","method","to","a","`","Model","`",".","Example",":","class","BaseModel",":","def","__init__","(","self",",","...",")",":","def","condition_on_observations","(","self",",","...",")",":","def","posterior","(","self",",","...",")",":","def","transform_inputs","(","self",",","...",")",":","class","ModelThatCanFantasize","(","BaseModel",",","FantasizeMixin",")",":","def","__init__","(","self",",","args",")",":","super","(",")",".__init__","(","args",")","model","=","ModelThatCanFantasize","(","...",")","model.fantasize","(","X",")"]
docstr: Construct a fantasy model.
  
  Constructs a fantasy model in the following fashion:
  (1) compute the model posterior at `X`, including observation noise.
  If `observation_noise` is a Tensor, use it directly as the observation
  noise to add.
  (2) sample from this posterior (using `sampler`) to generate "fake"
  observations.
  (3) condition the model on the new fake observations.
  
  Args:
      X: A `batch_shape x n' x d`-dim Tensor, where `d` is the dimension of
          the feature space, `n'` is the number of points per batch, and
          `batch_shape` is the batch shape (must be compatible with the
          batch shape of the model).
      sampler: The sampler used for sampling from the posterior at `X`.
      observation_noise: A `model_batch_shape x 1 x m`-dim tensor or
          a `model_batch_shape x n' x m`-dim tensor containing the average
          noise for each batch and output, where `m` is the number of outputs.
          `noise` must be in the outcome-transformed space if an outcome
          transform is used.
          If None and using an inferred noise likelihood, the noise will be the
          inferred noise level. If using a fixed noise likelihood, the mean across
          the observation noise in the training data is used as observation noise.
      kwargs: Will be passed to `model.condition_on_observations`
  
  Returns:
      The constructed fantasy model.
code: def fantasize(
      self: TFantasizeMixin,
      # TODO: see if any of these can be imported only if TYPE_CHECKING
      X: Tensor,
      sampler: MCSampler,
      observation_noise: Optional[Tensor] = None,
      **kwargs: Any,
  ) -> TFantasizeMixin:
      r"""Construct a fantasy model.
  
      Constructs a fantasy model in the following fashion:
      (1) compute the model posterior at `X`, including observation noise.
      If `observation_noise` is a Tensor, use it directly as the observation
      noise to add.
      (2) sample from this posterior (using `sampler`) to generate "fake"
      observations.
      (3) condition the model on the new fake observations.
  
      Args:
          X: A `batch_shape x n' x d`-dim Tensor, where `d` is the dimension of
              the feature space, `n'` is the number of points per batch, and
              `batch_shape` is the batch shape (must be compatible with the
              batch shape of the model).
          sampler: The sampler used for sampling from the posterior at `X`.
          observation_noise: A `model_batch_shape x 1 x m`-dim tensor or
              a `model_batch_shape x n' x m`-dim tensor containing the average
              noise for each batch and output, where `m` is the number of outputs.
              `noise` must be in the outcome-transformed space if an outcome
              transform is used.
              If None and using an inferred noise likelihood, the noise will be the
              inferred noise level. If using a fixed noise likelihood, the mean across
              the observation noise in the training data is used as observation noise.
          kwargs: Will be passed to `model.condition_on_observations`
  
      Returns:
          The constructed fantasy model.
      """
      if (
          not isinstance(observation_noise, Tensor)
          and observation_noise is not None
      ):
          raise DeprecationError(
              "`fantasize` no longer accepts a boolean for `observation_noise`."
          )
      elif observation_noise is None and isinstance(
          self.likelihood, FixedNoiseGaussianLikelihood
      ):
          if self.num_outputs > 1:
              # make noise ... x n x m
              observation_noise = self.likelihood.noise.transpose(
                  -1, -2
              )
          else:
              observation_noise = self.likelihood.noise.unsqueeze(-1)
          observation_noise = observation_noise.mean(
              dim=-2, keepdim=True
          )
      # if the inputs are empty, expand the inputs
      if X.shape[-2] == 0:
          output_shape = (
              sampler.sample_shape
              + X.shape[:-2]
              + self.batch_shape
              + torch.Size([0, self.num_outputs])
          )
          Y = torch.empty(output_shape, dtype=X.dtype, device=X.device)
          if observation_noise is not None:
              kwargs["noise"] = observation_noise.expand(Y.shape[1:])
          return self.condition_on_observations(
              X=self.transform_inputs(X),
              Y=Y,
              **kwargs,
          )
      propagate_grads = kwargs.pop("propagate_grads", False)
      with fantasize_flag():
          with settings.propagate_grads(propagate_grads):
              post_X = self.posterior(
                  X,
                  observation_noise=True
                  if observation_noise is None
                  else observation_noise,
              )
          Y_fantasized = sampler(
              post_X
          )  # num_fantasies x batch_shape x n' x m
          if observation_noise is not None:
              kwargs["noise"] = observation_noise.expand(
                  Y_fantasized.shape[1:]
              )
          return self.condition_on_observations(
              X=self.transform_inputs(X), Y=Y_fantasized, **kwargs
          )
  
imports: from __future__ import annotations
  import warnings
  from abc import ABC, abstractmethod
  from collections import defaultdict
  from copy import deepcopy
  from typing import Any, Callable, Dict, List, Mapping, Optional, Set, TYPE_CHECKING, TypeVar, Union
  import numpy
  import torch
  from botorch import settings
  from botorch.exceptions.errors import BotorchTensorDimensionError, DeprecationError, InputDataError
  from botorch.logging import shape_to_str
  from botorch.models.utils.assorted import fantasize
  from botorch.posteriors import Posterior, PosteriorList
  from botorch.sampling.base import MCSampler
  from botorch.sampling.list_sampler import ListSampler
  from botorch.utils.datasets import SupervisedDataset
  from botorch.utils.transforms import is_fully_bayesian
  from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood
  from torch import Tensor
  from torch.nn import Module, ModuleDict, ModuleList
label_desc: Use image node_id 4 for calling the FantasizeMixin obj's underlying member method code with example usage: obj.fantasize(X, sampler, observation_noise) and returns: self, self