code: def get_WPD_filter_v2(
      Phi: torch.Tensor,
      Rf: torch.Tensor,
      reference_vector: torch.Tensor,
      diagonal_loading: bool = True,
      diag_eps: float = 1e-7,
      eps: float = 1e-8,
  ) -> torch.Tensor:
      """Return the WPD vector (v2).
  
         This implementation is more efficient than `get_WPD_filter` as
          it skips unnecessary computation with zeros.
  
      Args:
          Phi (torch.complex64): (B, F, C, C)
              is speech PSD.
          Rf (torch.complex64): (B, F, (btaps+1) * C, (btaps+1) * C)
              is the power normalized spatio-temporal covariance matrix.
          reference_vector (torch.Tensor): (B, C)
              is the reference_vector.
          diagonal_loading (bool):
              Whether to add a tiny term to the diagonal of psd_n
          diag_eps (float):
          eps (float):
  
      Returns:
          filter_matrix (torch.complex64): (B, F, (btaps+1) * C)
      """
      C = reference_vector.shape[-1]
      if diagonal_loading:
          Rf = tik_reg(Rf, reg=diag_eps, eps=eps)
      inv_Rf = Rf.inverse()
      # (B, F, (btaps+1) * C, C)
      inv_Rf_pruned = inv_Rf[..., :C]
      # numerator: (..., C_1, C_2) x (..., C_2, C_3) -> (..., C_1, C_3)
      numerator = torch.matmul(inv_Rf_pruned, Phi)
      # NOTE (wangyou): until PyTorch 1.9.0, torch.trace does not
      # support bacth processing. Use FC.trace() as fallback.
      # ws: (..., (btaps+1) * C, C) / (...,) -> (..., (btaps+1) * C, C)
      ws = numerator / (
          FC.trace(numerator[..., :C, :])[..., None, None] + eps
      )
      # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
      beamform_vector = torch.einsum(
          "...fec,...c->...fe", ws, reference_vector.to(dtype=ws.dtype)
      )
      # (B, F, (btaps+1) * C)
      return beamform_vector
  
imports: from typing import List, Union
  import torch
  import torch_complex.functional
  import torchaudio
label_desc: Use image node_id 18 for calling a global function with example usage: get_WPD_filter_v2(Phi, Rf, reference_vector, diagonal_loading, diag_eps, eps) and returns: beamform_vector