docstr: Keep this API same with TasNet.
  
  Args:
      input: [M, K, N], M is batch size
      ilens (torch.Tensor): (M,)
      bottleneck_feat: [M, K, B]
      num_spk: number of speakers
      (Training: oracle,
      Inference: estimated by other module (e.g, EEND-EDA))
  
  Returns:
      masked (List[Union(torch.Tensor, ComplexTensor)]): [(M, K, N), ...]
      ilens (torch.Tensor): (M,)
      others predicted data, e.g. masks: OrderedDict[
          'mask_spk1': torch.Tensor(Batch, Frames, Freq),
          'mask_spk2': torch.Tensor(Batch, Frames, Freq),
          ...
          'mask_spkn': torch.Tensor(Batch, Frames, Freq),
      ]
code: def forward(
      self,
      input: Union[torch.Tensor, ComplexTensor],
      ilens: torch.Tensor,
      bottleneck_feat: torch.Tensor,
      num_spk: int,
  ) -> Tuple[
      List[Union[torch.Tensor, ComplexTensor]],
      torch.Tensor,
      OrderedDict,
  ]:
      """Keep this API same with TasNet.
  
      Args:
          input: [M, K, N], M is batch size
          ilens (torch.Tensor): (M,)
          bottleneck_feat: [M, K, B]
          num_spk: number of speakers
          (Training: oracle,
          Inference: estimated by other module (e.g, EEND-EDA))
  
      Returns:
          masked (List[Union(torch.Tensor, ComplexTensor)]): [(M, K, N), ...]
          ilens (torch.Tensor): (M,)
          others predicted data, e.g. masks: OrderedDict[
              'mask_spk1': torch.Tensor(Batch, Frames, Freq),
              'mask_spk2': torch.Tensor(Batch, Frames, Freq),
              ...
              'mask_spkn': torch.Tensor(Batch, Frames, Freq),
          ]
  
      """
      M, K, N = input.size()
      bottleneck_feat = bottleneck_feat.transpose(1, 2)  # [M, B, K]
      score = self.mask_conv1x1[num_spk - 1](
          bottleneck_feat
      )  # [M, B, K] -> [M, num_spk*N, K]
      # add other outputs of the module list with factor 0.0
      # to enable distributed training
      for z in range(self._max_num_spk):
          if z != num_spk - 1:
              score += 0.0 * F.interpolate(
                  self.mask_conv1x1[z](bottleneck_feat).transpose(1, 2),
                  size=num_spk * N,
              ).transpose(1, 2)
      score = score.view(
          M, num_spk, N, K
      )  # [M, num_spk*N, K] -> [M, num_spk, N, K]
      if self.mask_nonlinear == "softmax":
          est_mask = F.softmax(score, dim=1)
      elif self.mask_nonlinear == "relu":
          est_mask = F.relu(score)
      elif self.mask_nonlinear == "sigmoid":
          est_mask = torch.sigmoid(score)
      elif self.mask_nonlinear == "tanh":
          est_mask = torch.tanh(score)
      else:
          raise ValueError("Unsupported mask non-linear function")
  
      masks = est_mask.transpose(2, 3)  # [M, num_spk, K, N]
      masks = masks.unbind(dim=1)  # List[M, K, N]
  
      masked = [input * m for m in masks]
  
      others = OrderedDict(
          zip(
              ["mask_spk{}".format(i + 1) for i in range(len(masks))],
              masks,
          )
      )
  
      return masked, ilens, others
  
file_path: espnet/espnet2/diar/layers/multi_mask.py
imports: from collections import OrderedDict
  from typing import List, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from torch_complex.tensor import ComplexTensor
  from espnet2.diar.layers.abs_mask import AbsMask
label_desc: Use image node_id 3 for calling the MultiMask obj's underlying member method code with example usage: obj.forward(input, ilens, bottleneck_feat, num_spk) and returns: masked,  ilens,  others