code: def _update_omega(
      mean_f: Tensor,
      cov_f: Tensor,
      omega_f_nat_mean: Tensor,
      omega_f_nat_cov: Tensor,
      N: int,
      P: int,
      M: int,
      maximize: bool = True,
      jitter: float = 1e-6,
  ) -> Tuple[Tensor, Tensor]:
      r"""Computes the new omega factors by matching the moments.
  
      Args:
          mean_f: A `batch_shape x M x (N + P)`-dim Tensor containing the mean vector
              for the objectives.
          cov_f: A `batch_shape x M x (N + P) x (N + P)`-dim Tensor containing the
              covariance matrix for the objectives.
          omega_f_nat_mean: A `batch_shape x M x (N + P) x P x 2`-dim Tensor containing
              the omega natural mean factors for the objective matrix.
          omega_f_nat_cov: A `batch_shape x M x (N + P) x P x 2 x 2`-dim Tensor
              containing the omega natural covariance factors for the objective matrix.
          N: The number of design points.
          M: The number of Pareto optimal points.
          M: The number of objectives.
          maximize: If true, we consider the Pareto maximum domination relation.
          jitter: The jitter for the matrix inverse.
  
      Return:
          A two-element tuple containing
  
          - omega_f_nat_mean_new: A `batch_shape x M x (N + P) x P x 2` containing the
              new omega natural mean factors for the objective matrix.
          - omega_f_nat_cov_new: A `batch_shape x M x (N + P) x P x 2 x 2` containing
              the new omega natural covariance factors for the objective matrix.
      """
      tkwargs = {"dtype": mean_f.dtype, "device": mean_f.device}
      CLAMP_LB = torch.finfo(tkwargs["dtype"]).eps
      NEG_INF = torch.finfo(tkwargs["dtype"]).min
      weight = 1.0 if maximize else -1.0
      ###############################################################################
      # EXTRACT THE NECESSARY COMPONENTS
      ###############################################################################
      # `batch_shape x M x (N + P) x P x 2`-dim mean
      # `batch_shape x M x (N + P) x P x 2 x 2`-dim covariance
      mean_fX_fS, cov_fX_fS = _get_omega_f_contribution(
          mean_f, cov_f, N, P, M
      )
      identity = torch.diag_embed(
          torch.ones(cov_fX_fS.shape[:-1], **tkwargs)
      )
      # remove the Pareto diagonal
      cov_fX_fS = _replace_pareto_diagonal(
          cov_fX_fS + jitter * identity
      )
      nat_cov_fX_fS = torch.inverse(cov_fX_fS)
      nat_mean_fX_fS = torch.einsum(
          "...ij,...j->...i", nat_cov_fX_fS, mean_fX_fS
      )
  
      ###############################################################################
      # COMPUTE THE CAVITIES
      ###############################################################################
      # cavity distribution
      # natural parameters
      cav_nat_mean_f = nat_mean_fX_fS - omega_f_nat_mean
      cav_nat_cov_f = nat_cov_fX_fS - omega_f_nat_cov
  
      # transform to standard parameters
      # remove the Pareto diagonal
      cav_nat_cov_f = _replace_pareto_diagonal(cav_nat_cov_f)
      identity = torch.diag_embed(
          torch.ones(cav_nat_cov_f.shape[:-1], **tkwargs)
      )
      cav_cov_f = torch.inverse(cav_nat_cov_f + jitter * identity)
  
      cav_mean_f = torch.einsum(
          "...ij,...j->...i", cav_cov_f, cav_nat_mean_f
      )
  
      ###############################################################################
      # COMPUTE THE NORMALIZATION CONSTANT
      ###############################################################################
      # `batch_shape x M x (N + P) x P`
      # Equation 29
      cav_var_fX_minus_fS = (
          cav_cov_f[..., 0, 0]
          + cav_cov_f[..., 1, 1]
          - 2 * cav_cov_f[..., 0, 1]
      ).clamp_min(CLAMP_LB)
      cav_std_fX_minus_fS = torch.sqrt(cav_var_fX_minus_fS).clamp_min(
          CLAMP_LB
      )
  
      # `batch_shape x M x (N + P) x P`
      cav_mean_fX_minus_fS = weight * (
          cav_mean_f[..., 0] - cav_mean_f[..., 1]
      )
  
      # Equation 30
      cav_alpha = cav_mean_fX_minus_fS / cav_std_fX_minus_fS
      # compute alpha pdf and cdf
      normal_alpha = Normal(
          torch.zeros_like(cav_alpha), torch.ones_like(cav_alpha)
      )
      # `batch_shape x M x (N + P) x P`
      cav_alpha_log_cdf = log_cdf_robust(cav_alpha)
      # `batch_shape x M x (N + P) x P`
      cav_alpha_log_pdf = normal_alpha.log_prob(cav_alpha).clamp_min(
          NEG_INF
      )
      # `batch_shape x (N + P) x P`
      cav_sum_alpha_log_cdf = torch.sum(
          cav_alpha_log_cdf, dim=-3
      ).clamp_min(NEG_INF)
  
      # compute normalization constant Z
      # Equation 35
      cav_log_zeta = torch.log1p(
          -torch.exp(cav_sum_alpha_log_cdf)
      ).clamp_min(NEG_INF)
  
      # Need to clamp log values to prevent `exp(-inf) = nan`
      cav_logZ = cav_log_zeta
  
      # Equation 40 [first bit]
      # `batch_shape x (N + P) x P`
      cav_log_rho = -cav_logZ + cav_sum_alpha_log_cdf
  
      # Equation 40 [second bit]
      # `batch_shape x M x (N + P) x P`
      cav_log_rho = (
          cav_log_rho.unsqueeze(-3)
          - cav_alpha_log_cdf
          + cav_alpha_log_pdf
      )
      cav_rho = -torch.exp(cav_log_rho).clamp(NEG_INF, -NEG_INF)
      ###############################################################################
      # COMPUTE THE PARTIAL DERIVATIVES
      ###############################################################################
      # `batch_shape x M x (N + P) x P x 2`
      # Final vector: `[1, -1]`
      ones_mean = torch.ones(cav_mean_f.shape, **tkwargs)
      ones_mean[..., 1] = -ones_mean[..., 1]
  
      # `batch_shape x M x (N + P) x P x 2 x 2`
      # Final matrix: `[[1, -1], [-1, 1]]`
      ones_cov = torch.ones(cav_cov_f.shape, **tkwargs)
      ones_cov[..., 0, 1] = -ones_cov[..., 0, 1]
      ones_cov[..., 1, 0] = -ones_cov[..., 1, 0]
  
      # first partial derivation of the log Z with respect to the mean
      # assuming maximization (this is also where the sign will change)
      # Equation 41
      cav_dlogZ_dm = cav_rho / cav_std_fX_minus_fS
      cav_dlogZ_dm = weight * cav_dlogZ_dm.unsqueeze(-1) * ones_mean
  
      # second partial derivation of the log Z with respect to the mean
      # Equation 42
      cav_d2logZ_dm2 = (
          -cav_rho * (cav_rho + cav_alpha) / cav_var_fX_minus_fS
      )
      cav_d2logZ_dm2 = (
          cav_d2logZ_dm2.unsqueeze(-1).unsqueeze(-1) * ones_cov
      )
  
      ###############################################################################
      # COMPUTE THE NEW MEAN AND COVARIANCE
      ###############################################################################
      # compute the new mean and covariance
      cav_updated_mean_f = cav_mean_f + torch.einsum(
          "...ij,...j->...i", cav_cov_f, cav_dlogZ_dm
      )
      cav_updated_cov_f = cav_cov_f + torch.einsum(
          "...ij,...jk,...kl->...il",
          cav_cov_f,
          cav_d2logZ_dm2,
          cav_cov_f,
      )
      # transform to natural parameters
      # remove the Pareto diagonal
      cav_updated_cov_f = _replace_pareto_diagonal(cav_updated_cov_f)
  
      identity = torch.diag_embed(
          torch.ones(cav_updated_cov_f.shape[:-1], **tkwargs)
      )
      cav_updated_nat_cov_f = torch.inverse(
          cav_updated_cov_f + jitter * identity
      )
  
      cav_updated_nat_mean_f = torch.einsum(
          "...ij,...j->...i", cav_updated_nat_cov_f, cav_updated_mean_f
      )
  
      # match the moments to compute the gain
      omega_f_nat_mean_new = cav_updated_nat_mean_f - cav_nat_mean_f
      omega_f_nat_cov_new = cav_updated_nat_cov_f - cav_nat_cov_f
  
      # it is also possible to calculate the update directly as in the original paper:
      # identity = torch.diag_embed(torch.ones(cav_d2logZ_dm2.shape[:-1], **tkwargs))
      # denominator = torch.inverse(cav_cov_f @ cav_d2logZ_dm2 + identity)
      # omega_f_nat_cov_new = - cav_d2logZ_dm2 @ denominator
      # omega_f_nat_mean_new = torch.einsum(
      #     '...ij,...j->...i', denominator,
      #     cav_dlogZ_dm - torch.einsum('...ij,...j->...i', cav_d2logZ_dm2, cav_mean_f)
      # )
  
      return omega_f_nat_mean_new, omega_f_nat_cov_new
  
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 5 for calling a global function with example usage: _update_omega(mean_f, cov_f, omega_f_nat_mean, omega_f_nat_cov, N, P, M, maximize, jitter) and returns: omega_f_nat_mean_new,  omega_f_nat_cov_new