class_docstr: Estimate a reference channel by selecting the reference
  that maximizes the output SNR. It returns one-hot encoded
  vector or a soft reference.
  
  A straight-through estimator is used for gradient when using
  hard reference.
  
  Args:
      hard: If true, use hard estimate of ref channel.
          If false, use a soft estimate across channels.
      hard_use_grad: Use straight-through estimator for
          the gradient.
      subband_weighting: If true, use subband weighting when
          adding across subband SNRs. If false, use average
          across subbands.
  
  References:
      Boeddeker et al, Front-End Processing for the CHiME-5 Dinner Party Scenario, 2018
docstr: Args:
      W: Multichannel input multichannel output filter, shape (B, F, C, M), where
         C is the number of input channels and M is the number of output channels
      psd_s: Covariance for the signal, shape (B, F, C, C)
      psd_n: Covariance for the noise, shape (B, F, C, C)
  
  Returns:
      One-hot or soft reference channel, shape (B, M)
docstr_tok: ["Args",":","W",":","Multichannel","input","multichannel","output","filter",",","shape","(","B",",","F",",","C",",","M",")",",","where","C","is","the","number","of","input","channels","and","M","is","the","number","of","output","channels","psd_s",":","Covariance","for","the","signal",",","shape","(","B",",","F",",","C",",","C",")","psd_n",":","Covariance","for","the","noise",",","shape","(","B",",","F",",","C",",","C",")","Returns",":","One-hot","or","soft","reference","channel",",","shape","(","B",",","M",")"]
code: def forward(
      self, W: torch.Tensor, psd_s: torch.Tensor, psd_n: torch.Tensor
  ) -> torch.Tensor:
      """
      Args:
          W: Multichannel input multichannel output filter, shape (B, F, C, M), where
             C is the number of input channels and M is the number of output channels
          psd_s: Covariance for the signal, shape (B, F, C, C)
          psd_n: Covariance for the noise, shape (B, F, C, C)
  
      Returns:
          One-hot or soft reference channel, shape (B, M)
      """
      if self.subband_weighting:
          # (B, F, M)
          pow_s = torch.einsum(
              "...jm,...jk,...km->...m", W.conj(), psd_s, W
          ).abs()
          pow_n = torch.einsum(
              "...jm,...jk,...km->...m", W.conj(), psd_n, W
          ).abs()
  
          # Subband-weighting
          # (B, F, M) -> (B, M)
          pow_s = torch.sum(
              pow_s * self.weight_s.softmax(dim=0).unsqueeze(1), dim=-2
          )
          pow_n = torch.sum(
              pow_n * self.weight_n.softmax(dim=0).unsqueeze(1), dim=-2
          )
      else:
          # Sum across f as well
          # (B, F, C, M), (B, F, C, C), (B, F, C, M) -> (B, M)
          pow_s = torch.einsum(
              "...fjm,...fjk,...fkm->...m", W.conj(), psd_s, W
          ).abs()
          pow_n = torch.einsum(
              "...fjm,...fjk,...fkm->...m", W.conj(), psd_n, W
          ).abs()
  
      # Estimated SNR per channel (B, C)
      snr = pow_s / (pow_n + self.eps)
      snr = 10 * torch.log10(snr + self.eps)
  
      # Soft reference
      ref_soft = snr.softmax(dim=-1)
  
      if self.hard:
          _, idx = ref_soft.max(dim=-1, keepdim=True)
          ref_hard = torch.zeros_like(snr).scatter(-1, idx, 1.0)
          if self.hard_use_grad:
              # Straight-through for gradient
              # Propagate ref_soft gradient, as if thresholding is identity
              ref = ref_hard - ref_soft.detach() + ref_soft
          else:
              # No gradient
              ref = ref_hard
      else:
          ref = ref_soft
  
      return ref
  
imports: import random
  from typing import Callable, Optional
  import torch
  from nemo.collections.asr.parts.submodules.multi_head_attention import MultiHeadAttention
  from nemo.core.classes import NeuralModule, typecheck
  from nemo.core.neural_types import AudioSignal, FloatType, NeuralType, SpectrogramType
  from nemo.utils import logging
label_desc: Use image node_id 4 for calling the ReferenceChannelEstimatorSNR obj's underlying member method code with example usage: obj.forward(W, psd_s, psd_n) and returns: ref