code: def get_sdw_mwf_vector(
      psd_speech,
      psd_noise,
      reference_vector: Union[torch.Tensor, int],
      denoising_weight: float = 1.0,
      approx_low_rank_psd_speech: bool = False,
      iterations: int = 3,
      diagonal_loading: bool = True,
      diag_eps: float = 1e-7,
      eps: float = 1e-8,
  ):
      """Return the SDW-MWF (Speech Distortion Weighted Multi-channel Wiener Filter) vector
  
          h = (Spsd + mu * Npsd)^-1 @ Spsd @ u
  
      Reference:
          [1] Spatially pre-processed speech distortion weighted multi-channel Wiener
          filtering for noise reduction; A. Spriet et al, 2004
          https://dl.acm.org/doi/abs/10.1016/j.sigpro.2004.07.028
          [2] Rank-1 constrained multichannel Wiener filter for speech recognition in
          noisy environments; Z. Wang et al, 2018
          https://hal.inria.fr/hal-01634449/document
          [3] Low-rank approximation based multichannel Wiener filter algorithms for
          noise reduction with application in cochlear implants; R. Serizel, 2014
          https://ieeexplore.ieee.org/document/6730918
  
      Args:
          psd_speech (torch.complex64):
              speech covariance matrix (..., F, C, C)
          psd_noise (torch.complex64):
              noise covariance matrix (..., F, C, C)
          reference_vector (torch.Tensor or int): (..., C) or scalar
          denoising_weight (float): a trade-off parameter between noise reduction and
              speech distortion.
              A larger value leads to more noise reduction at the expense of more speech
              distortion.
              The plain MWF is obtained with `denoising_weight = 1` (by default).
          approx_low_rank_psd_speech (bool): whether to replace original input psd_speech
              with its low-rank approximation as in [2]
          iterations (int): number of iterations in power method, only used when
              `approx_low_rank_psd_speech = True`
          diagonal_loading (bool): Whether to add a tiny term to the diagonal of psd_n
          diag_eps (float):
          eps (float):
      Returns:
          beamform_vector (torch.complex64): (..., F, C)
      """  # noqa: H405, D205, D400, E501
      if approx_low_rank_psd_speech:
          if diagonal_loading:
              psd_noise = tik_reg(psd_noise, reg=diag_eps, eps=eps)
  
          # (B, F, C)
          recon_vec = get_rtf(
              psd_speech,
              psd_noise,
              mode="power",
              iterations=iterations,
              reference_vector=reference_vector,
              diagonal_loading=False,
          )
          # Eq. (25) in Ref[2]
          psd_speech_r1 = torch.einsum(
              "...c,...e->...ce", recon_vec, recon_vec.conj()
          )
          sigma_speech = FC.trace(psd_speech) / (
              FC.trace(psd_speech_r1) + eps
          )
          psd_speech_r1 = psd_speech_r1 * sigma_speech[..., None, None]
          # c.f. Eq. (62) in Ref[3]
          psd_speech = psd_speech_r1
  
      psd_n = psd_speech + denoising_weight * psd_noise
      if diagonal_loading:
          psd_n = tik_reg(psd_n, reg=diag_eps, eps=eps)
  
      ws = torch.linalg.solve(psd_n, psd_speech)
  
      if isinstance(reference_vector, int):
          beamform_vector = ws[..., reference_vector]
      else:
          beamform_vector = torch.einsum(
              "...fec,...c->...fe",
              ws,
              reference_vector.to(dtype=ws.dtype),
          )
      return beamform_vector
  
imports: from typing import List, Union
  import torch
  import torch_complex.functional
  import torchaudio
label_desc: Use image node_id 7 for calling a global function with example usage: get_sdw_mwf_vector(psd_speech, psd_noise, reference_vector, denoising_weight, approx_low_rank_psd_speech, iterations, diagonal_loading, diag_eps, eps) and returns: beamform_vector