class_docstr: The acquisition function for Max-value Entropy Search.
  
  This acquisition function computes the mutual information of max values and
  a candidate point X. See [Wang2017mves]_ for a detailed discussion.
  
  The model must be single-outcome. The batch case `q > 1` is supported
  through cyclic optimization and fantasies.
  
  Example:
      >>> model = SingleTaskGP(train_X, train_Y)
      >>> candidate_set = torch.rand(1000, bounds.size(1))
      >>> candidate_set = bounds[0] + (bounds[1] - bounds[0]) * candidate_set
      >>> MES = qMaxValueEntropy(model, candidate_set)
      >>> mes = MES(test_X)
docstr: Computes the information gain at the design points `X`.
  
  Approximately computes the information gain at the design points `X`,
  for both MES with noisy observations and multi-fidelity MES with noisy
  observation and trace observations.
  
  The implementation is inspired from the papers on multi-fidelity MES by
  [Takeno2020mfmves]_. The notation in the comments in this function follows
  the Appendix C of [Takeno2020mfmves]_.
  
  `num_fantasies = 1` for non-fantasized models.
  
  Args:
      X: A `batch_shape x 1 x d`-dim Tensor of `batch_shape` t-batches
          with `1` `d`-dim design point each.
      mean_M: A `batch_shape x num_fantasies x (m)`-dim Tensor of means.
      variance_M: A `batch_shape x num_fantasies x (m)`-dim Tensor of variances.
      covar_mM: A
          `batch_shape x num_fantasies x (m) x (1 + num_trace_observations)`-dim
          Tensor of covariances.
  
  Returns:
      A `num_fantasies x batch_shape`-dim Tensor of information gains at the
      given design points `X` (`num_fantasies=1` for non-fantasized models).
docstr_tok: ["Computes","the","information","gain","at","the","design","points","`","X","`",".","Approximately","computes","the","information","gain","at","the","design","points","`","X","`",",","for","both","MES","with","noisy","observations","and","multi-fidelity","MES","with","noisy","observation","and","trace","observations",".","The","implementation","is","inspired","from","the","papers","on","multi-fidelity","MES","by","[","Takeno2020mfmves","]","_",".","The","notation","in","the","comments","in","this","function","follows","the","Appendix","C","of","[","Takeno2020mfmves","]","_",".","`","num_fantasies","=","1","`","for","non-fantasized","models",".","Args",":","X",":","A","`","batch_shape","x","1","x","d","`","-dim","Tensor","of","`","batch_shape","`","t-batches","with","`","1","`","`","d","`","-dim","design","point","each",".","mean_M",":","A","`","batch_shape","x","num_fantasies","x","(","m",")","`","-dim","Tensor","of","means",".","variance_M",":","A","`","batch_shape","x","num_fantasies","x","(","m",")","`","-dim","Tensor","of","variances",".","covar_mM",":","A","`","batch_shape","x","num_fantasies","x","(","m",")","x","(","1","+","num_trace_observations",")","`","-dim","Tensor","of","covariances",".","Returns",":","A","`","num_fantasies","x","batch_shape","`","-dim","Tensor","of","information","gains","at","the","given","design","points","`","X","`","(","`","num_fantasies=1","`","for","non-fantasized","models",")","."]
code: def _compute_information_gain(
      self,
      X: Tensor,
      mean_M: Tensor,
      variance_M: Tensor,
      covar_mM: Tensor,
  ) -> Tensor:
      r"""Computes the information gain at the design points `X`.
  
      Approximately computes the information gain at the design points `X`,
      for both MES with noisy observations and multi-fidelity MES with noisy
      observation and trace observations.
  
      The implementation is inspired from the papers on multi-fidelity MES by
      [Takeno2020mfmves]_. The notation in the comments in this function follows
      the Appendix C of [Takeno2020mfmves]_.
  
      `num_fantasies = 1` for non-fantasized models.
  
      Args:
          X: A `batch_shape x 1 x d`-dim Tensor of `batch_shape` t-batches
              with `1` `d`-dim design point each.
          mean_M: A `batch_shape x num_fantasies x (m)`-dim Tensor of means.
          variance_M: A `batch_shape x num_fantasies x (m)`-dim Tensor of variances.
          covar_mM: A
              `batch_shape x num_fantasies x (m) x (1 + num_trace_observations)`-dim
              Tensor of covariances.
  
      Returns:
          A `num_fantasies x batch_shape`-dim Tensor of information gains at the
          given design points `X` (`num_fantasies=1` for non-fantasized models).
      """
      # compute the std_m, variance_m with noisy observation
      posterior_m = self.model.posterior(
          X.unsqueeze(-3),
          observation_noise=True,
          posterior_transform=self.posterior_transform,
      )
      # batch_shape x num_fantasies x (m) x (1 + num_trace_observations)
      mean_m = self.weight * posterior_m.mean.squeeze(-1)
      # batch_shape x num_fantasies x (m) x (1 + num_trace_observations)
      variance_m = posterior_m.distribution.covariance_matrix
      check_no_nans(variance_m)
  
      # compute mean and std for fM|ym, x, Dt ~ N(u, s^2)
      samples_m = self.weight * self.get_posterior_samples(
          posterior_m
      ).squeeze(-1)
      # s_m x batch_shape x num_fantasies x (m) (1 + num_trace_observations)
      L = psd_safe_cholesky(variance_m)
      temp_term = torch.cholesky_solve(
          covar_mM.unsqueeze(-1), L
      ).transpose(-2, -1)
      # equivalent to torch.matmul(covar_mM.unsqueeze(-2), torch.inverse(variance_m))
      # batch_shape x num_fantasies (m) x 1 x (1 + num_trace_observations)
  
      mean_pt1 = torch.matmul(
          temp_term, (samples_m - mean_m).unsqueeze(-1)
      )
      mean_new = mean_pt1.squeeze(-1).squeeze(-1) + mean_M
      # s_m x batch_shape x num_fantasies x (m)
      variance_pt1 = torch.matmul(temp_term, covar_mM.unsqueeze(-1))
      variance_new = variance_M - variance_pt1.squeeze(-1).squeeze(-1)
      # batch_shape x num_fantasies x (m)
      stdv_new = variance_new.clamp_min(CLAMP_LB).sqrt()
      # batch_shape x num_fantasies x (m)
  
      # define normal distribution to compute cdf and pdf
      normal = torch.distributions.Normal(
          torch.zeros(1, device=X.device, dtype=X.dtype),
          torch.ones(1, device=X.device, dtype=X.dtype),
      )
  
      # Compute p(fM <= f* | ym, x, Dt)
      view_shape = torch.Size(
          [
              self.posterior_max_values.shape[0],
              # add 1s to broadcast across the batch_shape of X
              *[
                  1
                  for _ in range(
                      X.ndim - self.posterior_max_values.ndim
                  )
              ],
              *self.posterior_max_values.shape[1:],
          ]
      )  # s_M x batch_shape x num_fantasies x (m)
      max_vals = self.posterior_max_values.view(view_shape).unsqueeze(1)
      # s_M x 1 x batch_shape x num_fantasies x (m)
      normalized_mvs_new = (max_vals - mean_new) / stdv_new
      # s_M x s_m x batch_shape x num_fantasies x (m)  =
      #   s_M x 1 x batch_shape x num_fantasies x (m)
      #   - s_m x batch_shape x num_fantasies x (m)
      cdf_mvs_new = normal.cdf(normalized_mvs_new).clamp_min(CLAMP_LB)
  
      # Compute p(fM <= f* | x, Dt)
      stdv_M = variance_M.sqrt()
      normalized_mvs = (max_vals - mean_M) / stdv_M
      # s_M x 1 x batch_shape x num_fantasies  x (m) =
      # s_M x 1 x 1 x num_fantasies x (m) - batch_shape x num_fantasies x (m)
      cdf_mvs = normal.cdf(normalized_mvs).clamp_min(CLAMP_LB)
      # s_M x 1 x batch_shape x num_fantasies x (m)
  
      # Compute log(p(ym | x, Dt))
      log_pdf_fm = posterior_m.distribution.log_prob(
          self.weight * samples_m
      ).unsqueeze(0)
      # 1 x s_m x batch_shape x num_fantasies x (m)
  
      # H0 = H(ym | x, Dt)
      H0 = (
          posterior_m.distribution.entropy()
      )  # batch_shape x num_fantasies x (m)
  
      # regression adjusted H1 estimation, H1_hat = H1_bar - beta * (H0_bar - H0)
      # H1 = E_{f*|x, Dt}[H(ym|f*, x, Dt)]
      Z = (
          cdf_mvs_new / cdf_mvs
      )  # s_M x s_m x batch_shape x num_fantasies x (m)
      # s_M x s_m x batch_shape x num_fantasies x (m)
      h1 = -Z * Z.log() - Z * log_pdf_fm
      check_no_nans(h1)
      dim = [0, 1]  # dimension of fm samples, fM samples
      H1_bar = h1.mean(dim=dim)
      h0 = -log_pdf_fm
      H0_bar = h0.mean(dim=dim)
      cov = ((h1 - H1_bar) * (h0 - H0_bar)).mean(dim=dim)
      beta = cov / (h0.var(dim=dim) * h1.var(dim=dim)).sqrt()
      H1_hat = H1_bar - beta * (H0_bar - H0)
      ig = H0 - H1_hat  # batch_shape x num_fantasies x (m)
      if self.posterior_max_values.ndim == 2:
          permute_idcs = [-1, *range(ig.ndim - 1)]
      else:
          permute_idcs = [-2, *range(ig.ndim - 2), -1]
      ig = ig.permute(
          *permute_idcs
      )  # num_fantasies x batch_shape x (m)
      return ig
  
imports: from __future__ import annotations
  from abc import ABC, abstractmethod
  from copy import deepcopy
  from math import log
  from typing import Any, Callable, Optional
  import numpy
  import torch
  from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
  from botorch.acquisition.cost_aware import CostAwareUtility, InverseCostWeightedUtility
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.exceptions.errors import UnsupportedError
  from botorch.models.cost import AffineFidelityCostModel
  from botorch.models.model import Model
  from botorch.models.utils import check_no_nans
  from botorch.sampling.normal import SobolQMCNormalSampler
  from botorch.utils.transforms import match_batch_shape, t_batch_mode_transform
  from linear_operator.functions import inv_quad
  from linear_operator.utils.cholesky import psd_safe_cholesky
  from scipy.optimize import brentq
  from scipy.stats import norm
  from torch import Tensor
label_desc: Use image node_id 3 for calling the qMaxValueEntropy obj's underlying member method code with example usage: obj._compute_information_gain(X, mean_M, variance_M, covar_mM) and returns: ig