code: def prepare_beamformer_stats(
      signal,
      masks_speech,
      mask_noise,
      powers=None,
      beamformer_type="mvdr",
      bdelay=3,
      btaps=5,
      eps=1e-6,
  ):
      """Prepare necessary statistics for constructing the specified beamformer.
  
      Args:
          signal (torch.complex64): (..., F, C, T)
          masks_speech (List[torch.Tensor]): (..., F, C, T) masks for all speech sources
          mask_noise (torch.Tensor): (..., F, C, T) noise mask
          powers (List[torch.Tensor]): powers for all speech sources (..., F, T)
                                       used for wMPDR or WPD beamformers
          beamformer_type (str): one of the pre-defined beamformer types
          bdelay (int): delay factor, used for WPD beamformser
          btaps (int): number of filter taps, used for WPD beamformser
          eps (torch.Tensor): tiny constant
      Returns:
          beamformer_stats (dict): a dictionary containing all necessary statistics
              e.g. "psd_n", "psd_speech", "psd_distortion"
              Note:
              * When `masks_speech` is a tensor or a single-element list, all returned
                statistics are tensors;
              * When `masks_speech` is a multi-element list, some returned statistics
                can be a list, e.g., "psd_n" for MVDR, "psd_speech" and "psd_distortion".
  
      """
      from espnet2.enh.layers.dnn_beamformer import BEAMFORMER_TYPES
  
      assert (
          beamformer_type in BEAMFORMER_TYPES
      ), "%s is not supported yet"
  
      if isinstance(masks_speech, (list, tuple)):
          masks_speech = [m.cdouble() for m in masks_speech]
      else:
          masks_speech = [masks_speech.cdouble()]
      num_spk = len(masks_speech)
  
      if (
          beamformer_type.startswith("wmpdr")
          or beamformer_type.startswith("wpd")
          or beamformer_type == "wlcmp"
          or beamformer_type == "wmwf"
      ):
          if powers is None:
              power_input = signal.real**2 + signal.imag**2
              # Averaging along the channel axis: (..., C, T) -> (..., T)
              powers = [
                  (power_input * m.abs()).mean(dim=-2)
                  for m in masks_speech
              ]
          else:
              assert len(powers) == num_spk, (len(powers), num_spk)
          inverse_powers = [1 / torch.clamp(p, min=eps) for p in powers]
  
      psd_transform = torchaudio.transforms.PSD(
          multi_mask=True, normalize=True
      )
      psd_speeches = [
          psd_transform(signal.transpose(-2, -3), m.transpose(-2, -3))
          for m in masks_speech
      ]
      if (
          beamformer_type == "mvdr_souden"
          or beamformer_type == "sdw_mwf"
          or beamformer_type == "r1mwf"
          or beamformer_type.startswith("mvdr_tfs")
          or not beamformer_type.endswith("_souden")
      ):
          # MVDR or other RTF-based formulas
          if mask_noise is not None:
              psd_bg = psd_transform(
                  signal.transpose(-2, -3),
                  mask_noise.cdouble().transpose(-2, -3),
              )
          if num_spk == 1:
              assert mask_noise is not None
              psd_noise = psd_bg
          else:
              psd_noise = []
              for i in range(num_spk):
                  if beamformer_type.startswith("mvdr_tfs"):
                      # NOTE: psd_noise is a list only for this beamformer
                      psd_noise_i = [
                          psd
                          for j, psd in enumerate(psd_speeches)
                          if j != i
                      ]
                  else:
                      psd_sum = sum(
                          psd
                          for j, psd in enumerate(psd_speeches)
                          if j != i
                      )
                      psd_noise_i = (
                          psd_bg + psd_sum
                          if mask_noise is not None
                          else psd_sum
                      )
                  psd_noise.append(psd_noise_i)
  
      if beamformer_type in (
          "mvdr",
          "mvdr_souden",
          "mvdr_tfs_souden",
          "sdw_mwf",
          "r1mwf",
          "lcmv",
          "gev",
          "gev_ban",
      ):
          psd_n = psd_noise
      elif beamformer_type == "mvdr_tfs":
          psd_n = psd_noise
          psd_noise = [sum(psd_noise_i) for psd_noise_i in psd_noise]
      elif beamformer_type in ("mpdr", "mpdr_souden", "lcmp", "mwf"):
          psd_n = torch.einsum(
              "...ct,...et->...ce", signal, signal.conj()
          )
      elif beamformer_type in (
          "wmpdr",
          "wmpdr_souden",
          "wlcmp",
          "wmwf",
      ):
          psd_n = [
              torch.einsum(
                  "...ct,...et->...ce",
                  signal * inv_p[..., None, :],
                  signal.conj(),
              )
              for inv_p in inverse_powers
          ]
      elif beamformer_type in ("wpd", "wpd_souden"):
          psd_n = [
              get_covariances(
                  signal, inv_p, bdelay, btaps, get_vector=False
              )
              for inv_p in inverse_powers
          ]
  
      if num_spk == 1:
          psd_speeches = psd_speeches[0]
          if isinstance(psd_n, (list, tuple)):
              psd_n = psd_n[0]
  
      if beamformer_type in (
          "mvdr",
          "mpdr",
          "wmpdr",
          "wpd",
          "lcmp",
          "wlcmp",
          "lcmv",
          "mvdr_tfs",
      ):
          return {
              "psd_n": psd_n,
              "psd_speech": psd_speeches,
              "psd_distortion": psd_noise,
          }
      elif (
          beamformer_type.endswith("_souden")
          or beamformer_type.startswith("gev")
          or beamformer_type == "mwf"
          or beamformer_type == "wmwf"
          or beamformer_type == "sdw_mwf"
          or beamformer_type == "r1mwf"
      ):
          return {"psd_n": psd_n, "psd_speech": psd_speeches}
  
imports: from typing import List, Union
  import torch
  import torch_complex.functional
  import torchaudio
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_beamformer_stats(signal, masks_speech, mask_noise, powers, beamformer_type, bdelay, btaps, eps) and returns: dict, dict