class_docstr: The acquisition function for Predictive Entropy Search. The code supports
  both single and multiple objectives as well as batching.
  
  This acquisition function approximates the mutual information between the
  observation at a candidate point `X` and the Pareto optimal input using the
  moment-matching procedure known as expectation propagation (EP).
  
  See the Appendix of [Garrido-Merchan2019]_ for the description of the EP
  procedure.
  
  IMPORTANT NOTES:
  (i) The PES acquisition function estimated using EP is sometimes not
  differentiable, and therefore we advise using a finite-difference estimate of
  the gradient as opposed to the gradients identified using automatic
  differentiation, which occasionally outputs `nan` values.
  
  The source of this differentiability is in the `_update_damping` function, which
  finds the damping factor `a` that is used to update the EP parameters
  `a * param_new + (1 - a) * param_old`. The damping factor has to ensure
  that the updated covariance matrices, `a * cov_f_new + (1 - a) cov_f_old`, is
  positive semi-definiteness. We follow the original paper, which identifies
  `a` via a successive halving scheme i.e. we check `a=1` then `a=0.5` etc. This
  procedure means `a` is a function of the test input `X`. This function is not
  differentiable  in `X`.
  
  (ii) EP could potentially fail for a number of reasons:
  
      (a) When the sampled Pareto optimal points `x_p` is poor compared to the
      training or testing data `x_n`.
  
      (b) When the training or testing data `x_n` is close the Pareto optimal
      points `x_p`.
  
      (c) When the convergence threshold is set too small.
  
  
      Problem (a) occurs because we have to compute the variable:
      `alpha = (mean(x_n) - mean(x_p)) / std(x_n - x_p)`, which becomes very
      large when `x_n` is better than `x_p` with high-probability. This leads to a
      log(0) error when we compute `log(1 - cdf(alpha))`. We have preemptively
      clamped some values depending on `1`alpha` in order to mitigate this.
  
      Problem (b) occurs because we have to compute matrix inverses for the
      two-dimensional marginals (x_n, x_p). To address this we manually add jitter
      to the diagonal of the covariance matrix i.e. `ep_jitter` when training and
      `test_jitter` when testing. The default choice is not always appropriate
      because the same jitter is used for the inversion of the covariance
      and precision matrix, which are on different scales.
  
      TODO: come up with strategy to adaptively update the jitter.
  
      Problem (c) occurs because a smaller threshold usually means that more EP
      iterations are required. Running too many EP iterations could lead to
      invertibility problems such as in problem (b). Setting a larger threshold
      or reducing the number of EP iterations could alleviate this.
  
  (iii) The estimated acquisition value could be negative.
class_docstr_tok: ["The","acquisition","function","for","Predictive","Entropy","Search",".","The","code","supports","both","single","and","multiple","objectives","as","well","as","batching",".","This","acquisition","function","approximates","the","mutual","information","between","the","observation","at","a","candidate","point","`","X","`","and","the","Pareto","optimal","input","using","the","moment-matching","procedure","known","as","expectation","propagation","(","EP",")",".","See","the","Appendix","of","[","Garrido-Merchan2019","]","_","for","the","description","of","the","EP","procedure",".","IMPORTANT","NOTES",":","(","i",")","The","PES","acquisition","function","estimated","using","EP","is","sometimes","not","differentiable",",","and","therefore","we","advise","using","a","finite-difference","estimate","of","the","gradient","as","opposed","to","the","gradients","identified","using","automatic","differentiation",",","which","occasionally","outputs","`","nan","`","values",".","The","source","of","this","differentiability","is","in","the","`","_update_damping","`","function",",","which","finds","the","damping","factor","`","a","`","that","is","used","to","update","the","EP","parameters","`","a","*","param_new","+","(","1","-","a",")","*","param_old","`",".","The","damping","factor","has","to","ensure","that","the","updated","covariance","matrices",",","`","a","*","cov_f_new","+","(","1","-","a",")","cov_f_old","`",",","is","positive","semi-definiteness",".","We","follow","the","original","paper",",","which","identifies","`","a","`","via","a","successive","halving","scheme","i.e",".","we","check","`","a=1","`","then","`","a=0.5","`","etc",".","This","procedure","means","`","a","`","is","a","function","of","the","test","input","`","X","`",".","This","function","is","not","differentiable","in","`","X","`",".","(","ii",")","EP","could","potentially","fail","for","a","number","of","reasons",":","(","a",")","When","the","sampled","Pareto","optimal","points","`","x_p","`","is","poor","compared","to","the","training","or","testing","data","`","x_n","`",".","(","b",")","When","the","training","or","testing","data","`","x_n","`","is","close","the","Pareto","optimal","points","`","x_p","`",".","(","c",")","When","the","convergence","threshold","is","set","too","small",".","Problem","(","a",")","occurs","because","we","have","to","compute","the","variable",":","`","alpha","=","(","mean","(","x_n",")","-","mean","(","x_p",")",")","\/","std","(","x_n","-","x_p",")","`",",","which","becomes","very","large","when","`","x_n","`","is","better","than","`","x_p","`","with","high-probability",".","This","leads","to","a","log","(","0",")","error","when","we","compute","`","log","(","1","-","cdf","(","alpha",")",")","`",".","We","have","preemptively","clamped","some","values","depending","on","`","1","`","alpha","`","in","order","to","mitigate","this",".","Problem","(","b",")","occurs","because","we","have","to","compute","matrix","inverses","for","the","two-dimensional","marginals","(","x_n",",","x_p",")",".","To","address","this","we","manually","add","jitter","to","the","diagonal","of","the","covariance","matrix","i.e",".","`","ep_jitter","`","when","training","and","`","test_jitter","`","when","testing",".","The","default","choice","is","not","always","appropriate","because","the","same","jitter","is","used","for","the","inversion","of","the","covariance","and","precision","matrix",",","which","are","on","different","scales",".","TODO",":","come","up","with","strategy","to","adaptively","update","the","jitter",".","Problem","(","c",")","occurs","because","a","smaller","threshold","usually","means","that","more","EP","iterations","are","required",".","Running","too","many","EP","iterations","could","lead","to","invertibility","problems","such","as","in","problem","(","b",")",".","Setting","a","larger","threshold","or","reducing","the","number","of","EP","iterations","could","alleviate","this",".","(","iii",")","The","estimated","acquisition","value","could","be","negative","."]
docstr: Perform expectation propagation to obtain the covariance factors that
  depend on the Pareto sets.
  
  The updates are performed in the natural parameter space. For a multivariate
  normal distribution with mean mu and covariance Sigma, we call Sigma^{-1}
  the natural covariance and Sigma^{-1} mu the natural mean.
docstr_tok: ["Perform","expectation","propagation","to","obtain","the","covariance","factors","that","depend","on","the","Pareto","sets",".","The","updates","are","performed","in","the","natural","parameter","space",".","For","a","multivariate","normal","distribution","with","mean","mu","and","covariance","Sigma",",","we","call","Sigma^","{","-1","}","the","natural","covariance","and","Sigma^","{","-1","}","mu","the","natural","mean","."]
code: def _expectation_propagation(self) -> None:
      r"""Perform expectation propagation to obtain the covariance factors that
      depend on the Pareto sets.
  
      The updates are performed in the natural parameter space. For a multivariate
      normal distribution with mean mu and covariance Sigma, we call Sigma^{-1}
      the natural covariance and Sigma^{-1} mu the natural mean.
      """
      ###########################################################################
      # INITIALIZATION
      ###########################################################################
      M = self.model.num_outputs
  
      if self.model.num_outputs > 1 or isinstance(
          self.model, ModelListGP
      ):
          train_X = self.model.train_inputs[0][0]
      else:
          train_X = self.model.train_inputs[0]
  
      tkwargs = {"dtype": train_X.dtype, "device": train_X.device}
      N = len(train_X)
      num_pareto_samples = self.num_pareto_samples
      P = self.pareto_sets.shape[-2]
  
      # initialize the predictive natural mean and variances
      (
          pred_nat_mean,
          pred_nat_cov,
          pred_mean,
          pred_cov,
      ) = _initialize_predictive_matrices(
          X=self.augmented_X,
          model=self.model,
          observation_noise=False,
          jitter=self.ep_jitter,
          natural=True,
      )
  
      pred_f_mean = pred_mean[..., 0:M, :]
      pred_f_nat_mean = pred_nat_mean[..., 0:M, :]
      pred_f_cov = pred_cov[..., 0:M, :, :]
      pred_f_nat_cov = pred_nat_cov[..., 0:M, :, :]
  
      # initialize the marginals
      # `num_pareto_samples x M x (N + P)`
      mean_f = pred_f_mean.clone()
      nat_mean_f = pred_f_nat_mean.clone()
      # `num_pareto_samples x M x (N + P) x (N + P)`
      cov_f = pred_f_cov.clone()
      nat_cov_f = pred_f_nat_cov.clone()
  
      # initialize omega the function which encodes the fact that the pareto points
      # are optimal in the feasible space i.e. any point in the feasible space
      # should not dominate the Pareto efficient points.
  
      # `num_pareto_samples x M x (N + P) x P x 2`
      omega_f_nat_mean = torch.zeros(
          (num_pareto_samples, M, N + P, P, 2), **tkwargs
      )
      # `num_pareto_samples x M x (N + P) x P x 2 x 2`
      omega_f_nat_cov = torch.zeros(
          (num_pareto_samples, M, N + P, P, 2, 2), **tkwargs
      )
  
      ###########################################################################
      # EXPECTATION PROPAGATION
      ###########################################################################
      damping = torch.ones(num_pareto_samples, M, **tkwargs)
  
      iteration = 0
      while (torch.sum(damping) > 0) and (
          iteration < self.max_ep_iterations
      ):
          # Compute the new natural mean and covariance
          ####################################################################
          # OBJECTIVE FUNCTION: OMEGA UPDATE
          ####################################################################
          (
              omega_f_nat_mean_new,
              omega_f_nat_cov_new,
          ) = _safe_update_omega(
              mean_f=mean_f,
              cov_f=cov_f,
              omega_f_nat_mean=omega_f_nat_mean,
              omega_f_nat_cov=omega_f_nat_cov,
              N=N,
              P=P,
              M=M,
              maximize=self.maximize,
              jitter=self.ep_jitter,
          )
  
          ####################################################################
          # OBJECTIVE FUNCTION: MARGINAL UPDATE
          ####################################################################
          nat_mean_f_new, nat_cov_f_new = _update_marginals(
              pred_f_nat_mean=pred_f_nat_mean,
              pred_f_nat_cov=pred_f_nat_cov,
              omega_f_nat_mean=omega_f_nat_mean_new,
              omega_f_nat_cov=omega_f_nat_cov_new,
              N=N,
              P=P,
          )
          ########################################################################
          # OBJECTIVE FUNCTION: DAMPING UPDATE
          ########################################################################
          # update damping of objectives
          damping, cholesky_nat_cov_f = _update_damping(
              nat_cov=nat_cov_f,
              nat_cov_new=nat_cov_f_new,
              damping_factor=damping,
              jitter=self.ep_jitter,
          )
          check_no_nans(cholesky_nat_cov_f)
          ########################################################################
          # OBJECTIVE FUNCTION: DAMPED UPDATE
          ########################################################################
          # Damp update of omega
          omega_f_nat_mean = _damped_update(
              old_factor=omega_f_nat_mean,
              new_factor=omega_f_nat_mean_new,
              damping_factor=damping,
          )
  
          omega_f_nat_cov = _damped_update(
              old_factor=omega_f_nat_cov,
              new_factor=omega_f_nat_cov_new,
              damping_factor=damping,
          )
          # update the mean and covariance
          nat_mean_f = _damped_update(
              old_factor=nat_mean_f,
              new_factor=nat_mean_f_new,
              damping_factor=damping,
          )
          nat_cov_f = _damped_update(
              old_factor=nat_cov_f,
              new_factor=nat_cov_f_new,
              damping_factor=damping,
          )
  
          # compute cholesky inverse
          cov_f_new = torch.cholesky_inverse(cholesky_nat_cov_f)
          mean_f_new = torch.einsum(
              "...ij,...j->...i", cov_f_new, nat_mean_f
          )
          check_no_nans(cov_f_new)
          ########################################################################
          # OBJECTIVE FUNCTION: CONVERGENCE UPDATE
          ########################################################################
          # Set the damping to zero when the change in the mean and
          # covariance is less than the threshold
          (
              damping,
              delta_mean_f,
              delta_cov_f,
          ) = _update_damping_when_converged(
              mean_old=mean_f,
              mean_new=mean_f_new,
              cov_old=cov_f,
              cov_new=cov_f_new,
              damping_factor=damping,
              threshold=self.threshold,
              iteration=iteration,
          )
          cov_f = cov_f_new
          mean_f = mean_f_new
          iteration = iteration + 1
  
      ############################################################################
      # SAVE OMEGA AND PHI FACTORS
      ############################################################################
      check_no_nans(omega_f_nat_mean)
      check_no_nans(omega_f_nat_cov)
      # save phi and omega for the forward
      self._omega_f_nat_mean = omega_f_nat_mean
      self._omega_f_nat_cov = omega_f_nat_cov
  
imports: from __future__ import annotations
  from typing import Any, Optional, Tuple
  import torch
  from botorch.acquisition.acquisition import AcquisitionFunction
  from botorch.exceptions import InputDataError
  from botorch.exceptions.errors import UnsupportedError
  from botorch.models.model import Model
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.models.utils import check_no_nans
  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 2 for calling the qMultiObjectivePredictiveEntropySearch obj's underlying member method code with example usage: obj._expectation_propagation() without return types