code: def _safe_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"""Try to update the new omega factors by matching the moments. If the update
      is not possible then this returns the initial omega factors.
  
      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.
      """
      try:
          omega_f_nat_mean_new, omega_f_nat_cov_new = _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=maximize,
              jitter=jitter,
          )
          check_no_nans(omega_f_nat_mean_new)
          check_no_nans(omega_f_nat_cov_new)
          return omega_f_nat_mean_new, omega_f_nat_cov_new
  
      except RuntimeError or InputDataError:
          return omega_f_nat_mean, 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 6 for calling a global function with example usage: _safe_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, omega_f_nat_mean,  omega_f_nat_cov