docstr: Forward.
  
  Args:
      input (torch.complex64/ComplexTensor):
          mixed speech [Batch, Frames, Channel, Freq]
      ilens (torch.Tensor): input lengths [Batch]
      additional (Dict or None): other data included in model
          NOTE: not used in this model
  
  Returns:
      enhanced speech (single-channel): List[torch.complex64/ComplexTensor]
      output lengths
      other predcited data: OrderedDict[
          'dereverb1': ComplexTensor(Batch, Frames, Channel, Freq),
          'mask_dereverb1': torch.Tensor(Batch, Frames, Channel, Freq),
          'mask_noise1': torch.Tensor(Batch, Frames, Channel, Freq),
          'mask_spk1': torch.Tensor(Batch, Frames, Channel, Freq),
          'mask_spk2': torch.Tensor(Batch, Frames, Channel, Freq),
          ...
          'mask_spkn': torch.Tensor(Batch, Frames, Channel, Freq),
      ]
code: def forward(
      self,
      input: Union[torch.Tensor, ComplexTensor],
      ilens: torch.Tensor,
      additional: Optional[Dict] = None,
  ) -> Tuple[
      List[Union[torch.Tensor, ComplexTensor]],
      torch.Tensor,
      OrderedDict,
  ]:
      """Forward.
  
      Args:
          input (torch.complex64/ComplexTensor):
              mixed speech [Batch, Frames, Channel, Freq]
          ilens (torch.Tensor): input lengths [Batch]
          additional (Dict or None): other data included in model
              NOTE: not used in this model
  
      Returns:
          enhanced speech (single-channel): List[torch.complex64/ComplexTensor]
          output lengths
          other predcited data: OrderedDict[
              'dereverb1': ComplexTensor(Batch, Frames, Channel, Freq),
              'mask_dereverb1': torch.Tensor(Batch, Frames, Channel, Freq),
              'mask_noise1': torch.Tensor(Batch, Frames, Channel, Freq),
              'mask_spk1': torch.Tensor(Batch, Frames, Channel, Freq),
              'mask_spk2': torch.Tensor(Batch, Frames, Channel, Freq),
              ...
              'mask_spkn': torch.Tensor(Batch, Frames, Channel, Freq),
          ]
      """
      # Shape of input spectrum must be (B, T, F) or (B, T, C, F)
      assert input.dim() in (3, 4), input.dim()
      enhanced = input
      others = OrderedDict()
  
      if (
          self.training
          and self.loss_type is not None
          and self.loss_type.startswith("mask")
      ):
          # Only estimating masks during training for saving memory
          if self.use_wpe:
              if input.dim() == 3:
                  mask_w, ilens = self.wpe.predict_mask(
                      input.unsqueeze(-2), ilens
                  )
                  mask_w = mask_w.squeeze(-2)
              elif input.dim() == 4:
                  mask_w, ilens = self.wpe.predict_mask(input, ilens)
  
              if mask_w is not None:
                  if isinstance(enhanced, list):
                      # single-source WPE
                      for spk in range(self.num_spk):
                          others[
                              "mask_dereverb{}".format(spk + 1)
                          ] = mask_w[spk]
                  else:
                      # multi-source WPE
                      others["mask_dereverb1"] = mask_w
  
          if self.use_beamformer and input.dim() == 4:
              others_b, ilens = self.beamformer.predict_mask(
                  input, ilens
              )
              for spk in range(self.num_spk):
                  others["mask_spk{}".format(spk + 1)] = others_b[spk]
              if len(others_b) > self.num_spk:
                  others["mask_noise1"] = others_b[self.num_spk]
  
          return None, ilens, others
  
      else:
          powers = None
          # Performing both mask estimation and enhancement
          if input.dim() == 3:
              # single-channel input (B, T, F)
              if self.use_wpe:
                  enhanced, ilens, mask_w, powers = self.wpe(
                      input.unsqueeze(-2), ilens
                  )
                  if isinstance(enhanced, list):
                      # single-source WPE
                      enhanced = [enh.squeeze(-2) for enh in enhanced]
                      if mask_w is not None:
                          for spk in range(self.num_spk):
                              key = "dereverb{}".format(spk + 1)
                              others[key] = enhanced[spk]
                              others["mask_" + key] = mask_w[
                                  spk
                              ].squeeze(-2)
                  else:
                      # multi-source WPE
                      enhanced = enhanced.squeeze(-2)
                      if mask_w is not None:
                          others["dereverb1"] = enhanced
                          others["mask_dereverb1"] = mask_w.squeeze(-2)
          else:
              # multi-channel input (B, T, C, F)
              # 1. WPE
              if self.use_wpe:
                  enhanced, ilens, mask_w, powers = self.wpe(
                      input, ilens
                  )
                  if mask_w is not None:
                      if isinstance(enhanced, list):
                          # single-source WPE
                          for spk in range(self.num_spk):
                              key = "dereverb{}".format(spk + 1)
                              others[key] = enhanced[spk]
                              others["mask_" + key] = mask_w[spk]
                      else:
                          # multi-source WPE
                          others["dereverb1"] = enhanced
                          others["mask_dereverb1"] = mask_w.squeeze(-2)
  
              # 2. Beamformer
              if self.use_beamformer:
                  if (
                      not self.beamformer.beamformer_type.startswith(
                          "wmpdr"
                      )
                      or not self.beamformer.beamformer_type.startswith(
                          "wpd"
                      )
                      or not self.shared_power
                      or (self.wpe.nmask == 1 and self.num_spk > 1)
                  ):
                      powers = None
  
                  # enhanced: (B, T, C, F) -> (B, T, F)
                  if isinstance(enhanced, list):
                      # outputs of single-source WPE
                      raise NotImplementedError(
                          "Single-source WPE is not supported with beamformer "
                          "in multi-speaker cases."
                      )
                  else:
                      # output of multi-source WPE
                      enhanced, ilens, others_b = self.beamformer(
                          enhanced, ilens, powers=powers
                      )
                  for spk in range(self.num_spk):
                      others["mask_spk{}".format(spk + 1)] = others_b[
                          spk
                      ]
                  if len(others_b) > self.num_spk:
                      others["mask_noise1"] = others_b[self.num_spk]
  
      if not isinstance(enhanced, list):
          enhanced = [enhanced]
  
      return enhanced, ilens, others
  
file_path: espnet/espnet2/enh/separator/neural_beamformer.py
imports: from collections import OrderedDict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from torch_complex.tensor import ComplexTensor
  from espnet2.enh.layers.dnn_beamformer import DNN_Beamformer
  from espnet2.enh.layers.dnn_wpe import DNN_WPE
  from espnet2.enh.separator.abs_separator import AbsSeparator
label_desc: Use image node_id 2 for calling the NeuralBeamformer obj's underlying member method code with example usage: obj.forward(input, ilens, additional) and returns: enhanced,  ilens,  others, None,  ilens,  others