class_docstr: The acquisition function for the Joint Entropy Search, where the batches
  `q > 1` are supported through the lower bound formulation.
  
  This acquisition function computes the mutual information between the observation
  at a candidate point `X` and the optimal input-output pair.
  
  See [Tu2022joint]_ for a discussion on the estimation procedure.
class_docstr_tok: ["The","acquisition","function","for","the","Joint","Entropy","Search",",","where","the","batches","`","q",">","1","`","are","supported","through","the","lower","bound","formulation",".","This","acquisition","function","computes","the","mutual","information","between","the","observation","at","a","candidate","point","`","X","`","and","the","optimal","input-output","pair",".","See","[","Tu2022joint","]","_","for","a","discussion","on","the","estimation","procedure","."]
docstr: Joint entropy search acquisition function.
  
  Args:
      model: A fitted single-outcome model.
      optimal_inputs: A `num_samples x d`-dim tensor containing the sampled
          optimal inputs of dimension `d`. We assume for simplicity that each
          sample only contains one optimal set of inputs.
      optimal_outputs: A `num_samples x 1`-dim Tensor containing the optimal
          set of objectives of dimension `1`.
      condition_noiseless: Whether to condition on noiseless optimal observations
          `f*` [Hvarfner2022joint]_ or noisy optimal observations `y*`
          [Tu2022joint]_. These are sampled identically, so this only controls
          the fashion in which the GP is reshaped as a result of conditioning
          on the optimum.
      estimation_type: estimation_type: A string to determine which entropy
          estimate is computed: Lower bound" ("LB") or "Monte Carlo" ("MC").
          Lower Bound is recommended due to the relatively high variance
          of the MC estimator.
      maximize: If true, we consider a maximization problem.
      X_pending: A `m x d`-dim Tensor of `m` design points that have been
          submitted for function evaluation, but have not yet been evaluated.
      num_samples: The number of Monte Carlo samples used for the Monte Carlo
          estimate.
docstr_tok: ["Joint","entropy","search","acquisition","function",".","Args",":","model",":","A","fitted","single-outcome","model",".","optimal_inputs",":","A","`","num_samples","x","d","`","-dim","tensor","containing","the","sampled","optimal","inputs","of","dimension","`","d","`",".","We","assume","for","simplicity","that","each","sample","only","contains","one","optimal","set","of","inputs",".","optimal_outputs",":","A","`","num_samples","x","1","`","-dim","Tensor","containing","the","optimal","set","of","objectives","of","dimension","`","1","`",".","condition_noiseless",":","Whether","to","condition","on","noiseless","optimal","observations","`","f","*","`","[","Hvarfner2022joint","]","_","or","noisy","optimal","observations","`","y","*","`","[","Tu2022joint","]","_",".","These","are","sampled","identically",",","so","this","only","controls","the","fashion","in","which","the","GP","is","reshaped","as","a","result","of","conditioning","on","the","optimum",".","estimation_type",":","estimation_type",":","A","string","to","determine","which","entropy","estimate","is","computed",":","Lower","bound","''","(","``","LB","''",")","or","``","Monte","Carlo","''","(","``","MC","''",")",".","Lower","Bound","is","recommended","due","to","the","relatively","high","variance","of","the","MC","estimator",".","maximize",":","If","true",",","we","consider","a","maximization","problem",".","X_pending",":","A","`","m","x","d","`","-dim","Tensor","of","`","m","`","design","points","that","have","been","submitted","for","function","evaluation",",","but","have","not","yet","been","evaluated",".","num_samples",":","The","number","of","Monte","Carlo","samples","used","for","the","Monte","Carlo","estimate","."]
code: def __init__(
      self,
      model: Model,
      optimal_inputs: Tensor,
      optimal_outputs: Tensor,
      condition_noiseless: bool = True,
      posterior_transform: Optional[PosteriorTransform] = None,
      X_pending: Optional[Tensor] = None,
      estimation_type: str = "LB",
      maximize: bool = True,
      num_samples: int = 64,
  ) -> None:
      r"""Joint entropy search acquisition function.
  
      Args:
          model: A fitted single-outcome model.
          optimal_inputs: A `num_samples x d`-dim tensor containing the sampled
              optimal inputs of dimension `d`. We assume for simplicity that each
              sample only contains one optimal set of inputs.
          optimal_outputs: A `num_samples x 1`-dim Tensor containing the optimal
              set of objectives of dimension `1`.
          condition_noiseless: Whether to condition on noiseless optimal observations
              `f*` [Hvarfner2022joint]_ or noisy optimal observations `y*`
              [Tu2022joint]_. These are sampled identically, so this only controls
              the fashion in which the GP is reshaped as a result of conditioning
              on the optimum.
          estimation_type: estimation_type: A string to determine which entropy
              estimate is computed: Lower bound" ("LB") or "Monte Carlo" ("MC").
              Lower Bound is recommended due to the relatively high variance
              of the MC estimator.
          maximize: If true, we consider a maximization problem.
          X_pending: A `m x d`-dim Tensor of `m` design points that have been
              submitted for function evaluation, but have not yet been evaluated.
          num_samples: The number of Monte Carlo samples used for the Monte Carlo
              estimate.
      """
      super().__init__(model=model)
      sampler = SobolQMCNormalSampler(
          sample_shape=torch.Size([num_samples])
      )
      MCSamplerMixin.__init__(self, sampler=sampler)
      # To enable fully bayesian GP conditioning, we need to unsqueeze
      # to get num_optima x num_gps unique GPs
  
      # inputs come as num_optima_per_model x (num_models) x d
      # but we want it four-dimensional in the Fully bayesian case,
      # and three-dimensional otherwise.
      self.optimal_inputs = optimal_inputs.unsqueeze(-2)
      self.optimal_outputs = optimal_outputs.unsqueeze(-2)
      self.posterior_transform = posterior_transform
      self.maximize = maximize
  
      # The optima (can be maxima, can be minima) come in as the largest
      # values if we optimize, or the smallest (likely substantially negative)
      # if we minimize. Inside the acquisition function, however, we always
      # want to consider MAX-values. As such, we need to flip them if
      # we want to minimize.
      if not self.maximize:
          optimal_outputs = -optimal_outputs
      self.num_samples = optimal_inputs.shape[0]
      self.condition_noiseless = condition_noiseless
      self.initial_model = model
  
      # Here, the optimal inputs have shapes num_optima x [num_models if FB] x 1 x D
      # and the optimal outputs have shapes num_optima x [num_models if FB] x 1 x 1
      # The third dimension equaling 1 is required to get one optimum per model,
      # which raises a BotorchTensorDimensionWarning.
      if isinstance(model, SaasFullyBayesianSingleTaskGP):
          raise NotImplementedError(FULLY_BAYESIAN_ERROR_MSG)
      with warnings.catch_warnings():
          warnings.filterwarnings("ignore")
          with fantasize_flag():
              with settings.propagate_grads(False):
                  # We must do a forward pass one before conditioning.
                  self.initial_model.posterior(
                      self.optimal_inputs[:1], observation_noise=False
                  )
  
              # This equates to the JES version proposed by Hvarfner et. al.
              if self.condition_noiseless:
                  opt_noise = torch.full_like(
                      self.optimal_outputs, MIN_INFERRED_NOISE_LEVEL
                  )
                  # conditional (batch) model of shape (num_models)
                  # x num_optima_per_model
                  self.conditional_model = (
                      self.initial_model.condition_on_observations(
                          X=self.initial_model.transform_inputs(
                              self.optimal_inputs
                          ),
                          Y=self.optimal_outputs,
                          noise=opt_noise,
                      )
                  )
              else:
                  self.conditional_model = (
                      self.initial_model.condition_on_observations(
                          X=self.initial_model.transform_inputs(
                              self.optimal_inputs
                          ),
                          Y=self.optimal_outputs,
                      )
                  )
  
      self.estimation_type = estimation_type
      self.set_X_pending(X_pending)
  
imports: from __future__ import annotations
  import warnings
  from math import log, pi
  from typing import Optional
  import torch
  from botorch import settings
  from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
  from botorch.models.gp_regression import MIN_INFERRED_NOISE_LEVEL
  from botorch.models.model import Model
  from botorch.models.utils import check_no_nans, fantasize
  from botorch.sampling.normal import SobolQMCNormalSampler
  from botorch.utils.transforms import concatenate_pending_points, t_batch_mode_transform
  from torch import Tensor
  from torch.distributions import Normal
label_desc: Use image node_id 1 to create a new qJointEntropySearch object from inherited base classes: AcquisitionFunction, MCSamplerMixin with example: obj = qJointEntropySearch(model, optimal_inputs, optimal_outputs, condition_noiseless, posterior_transform, X_pending, estimation_type, maximize, num_samples)