code: def _compute_entropy_upper_bound(
      hypercell_bounds: Tensor,
      mean: Tensor,
      variance: Tensor,
      observation_noise: Tensor,
      only_diagonal: bool = False,
  ) -> Tensor:
      r"""Computes the entropy upper bound at the design points `X`. This is used for
      the JES-LB and MES-LB estimate. If `only_diagonal` is True, then this computes
      the entropy estimate for the JES-LB2 and MES-LB2.
  
      Args:
          hypercell_bounds: A `num_pareto_samples x 2 x J x M` -dim Tensor containing
              the box decomposition bounds, where `J` = max(num_boxes).
          mean: A `batch_shape x num_pareto_samples x q x 1 x M`-dim Tensor containing
              the posterior mean at X.
          variance: A `batch_shape x num_pareto_samples x q x 1 x M`-dim Tensor
              containing the posterior variance at X excluding observation noise.
          observation_noise: A `batch_shape x num_pareto_samples x q x 1 x M`-dim
              Tensor containing the observation noise at X.
          only_diagonal: If true, we only compute the diagonal elements of the variance.
  
      Returns:
          A `batch_shape x q`-dim Tensor of entropy estimate at the given design points
          `X`.
      """
      tkwargs = {
          "dtype": hypercell_bounds.dtype,
          "device": hypercell_bounds.device,
      }
      CLAMP_LB = torch.finfo(tkwargs["dtype"]).eps
  
      variance_plus_noise = variance + observation_noise
  
      # Standardize the box decomposition bounds and compute normal quantities.
      # `batch_shape x num_pareto_samples x q x 2 x J x M`
      g = (
          hypercell_bounds.unsqueeze(-4) - mean.unsqueeze(-2)
      ) / torch.sqrt(variance.unsqueeze(-2))
      normal = Normal(torch.zeros_like(g), torch.ones_like(g))
      gcdf = normal.cdf(g)
      gpdf = torch.exp(normal.log_prob(g))
      g_times_gpdf = g * gpdf
  
      # Compute the differences between the upper and lower terms.
      Wjm = (gcdf[..., 1, :, :] - gcdf[..., 0, :, :]).clamp_min(
          CLAMP_LB
      )
      Vjm = g_times_gpdf[..., 1, :, :] - g_times_gpdf[..., 0, :, :]
      Gjm = gpdf[..., 1, :, :] - gpdf[..., 0, :, :]
  
      # Compute W.
      Wj = torch.exp(torch.sum(torch.log(Wjm), dim=-1, keepdims=True))
      W = torch.sum(Wj, dim=-2, keepdims=True).clamp_max(1.0)
  
      Cjm = Gjm / Wjm
  
      # First moment:
      Rjm = Cjm * Wj / W
      # `batch_shape x num_pareto_samples x q x 1 x M
      mom1 = mean - torch.sqrt(variance) * Rjm.sum(-2, keepdims=True)
      # diagonal weighted sum
      # `batch_shape x num_pareto_samples x q x 1 x M
      diag_weighted_sum = (Wj * variance * Vjm / Wjm / W).sum(
          -2, keepdims=True
      )
  
      if only_diagonal:
          # `batch_shape x num_pareto_samples x q x 1 x M`
          mean_squared = mean.pow(2)
          cross_sum = -2 * (mean * torch.sqrt(variance) * Rjm).sum(
              -2, keepdims=True
          )
          # `batch_shape x num_pareto_samples x q x 1 x M`
          mom2 = (
              variance_plus_noise
              - diag_weighted_sum
              + cross_sum
              + mean_squared
          )
          var = (mom2 - mom1.pow(2)).clamp_min(CLAMP_LB)
  
          # `batch_shape x num_pareto_samples x q
          log_det_term = 0.5 * torch.log(var).sum(dim=-1).squeeze(-1)
      else:
          # First moment x First moment
          # `batch_shape x num_pareto_samples x q x 1 x M x M
          cross_mom1 = torch.einsum("...i,...j->...ij", mom1, mom1)
  
          # Second moment:
          # `batch_shape x num_pareto_samples x q x 1 x M x M
          # firstly compute the general terms
          mom2_cross1 = -torch.einsum(
              "...i,...j->...ij", mean, torch.sqrt(variance) * Cjm
          )
          mom2_cross2 = -torch.einsum(
              "...i,...j->...ji", mean, torch.sqrt(variance) * Cjm
          )
          mom2_mean_squared = torch.einsum(
              "...i,...j->...ij", mean, mean
          )
  
          mom2_weighted_sum = (
              (mom2_cross1 + mom2_cross2)
              * Wj.unsqueeze(-1)
              / W.unsqueeze(-1)
          ).sum(-3, keepdims=True)
          mom2_weighted_sum = mom2_weighted_sum + mom2_mean_squared
  
          # Compute the additional off-diagonal terms.
          mom2_off_diag = torch.einsum(
              "...i,...j->...ij",
              torch.sqrt(variance) * Cjm,
              torch.sqrt(variance) * Cjm,
          )
          mom2_off_diag_sum = (
              mom2_off_diag * Wj.unsqueeze(-1) / W.unsqueeze(-1)
          ).sum(-3, keepdims=True)
  
          # Compute the diagonal terms and subtract the diagonal computed before.
          init_diag = torch.diagonal(
              mom2_off_diag_sum, dim1=-2, dim2=-1
          )
          diag_weighted_sum = torch.diag_embed(
              variance_plus_noise - diag_weighted_sum - init_diag
          )
          mom2 = (
              mom2_weighted_sum + mom2_off_diag_sum + diag_weighted_sum
          )
          # Compute the variance
          var = (mom2 - cross_mom1).squeeze(-3)
  
          # Jitter the diagonal.
          # The jitter is probably not needed here at all.
          jitter_diag = 1e-6 * torch.diag_embed(
              torch.ones(var.shape[:-1], **tkwargs)
          )
          log_det_term = 0.5 * torch.logdet(var + jitter_diag)
  
      # Additional terms.
      M_plus_K = mean.shape[-1]
      add_term = (
          0.5
          * M_plus_K
          * (1 + torch.log(torch.ones(1, **tkwargs) * 2 * pi))
      )
  
      # `batch_shape x num_pareto_samples x q
      entropy = add_term + log_det_term
      return entropy.mean(-2)
  
imports: from __future__ import annotations
  from abc import abstractmethod
  from math import pi
  from typing import Any, Optional, Tuple, Union
  import torch
  from botorch import settings
  from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
  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 fantasize
  from botorch.posteriors.gpytorch import GPyTorchPosterior
  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 2 for calling a global function with example usage: _compute_entropy_upper_bound(hypercell_bounds, mean, variance, observation_noise, only_diagonal) and returns: entropy