class_docstr: Target Speaker Extraction Frontend model
class_docstr_tok: ["Target","Speaker","Extraction","Frontend","model"]
docstr: Frontend + Encoder + Decoder + Calc loss
  
  Args:
      speech_mix: (Batch, samples) or (Batch, samples, channels)
      speech_ref1: (Batch, samples)
                  or (Batch, samples, channels)
      speech_ref2: (Batch, samples)
                  or (Batch, samples, channels)
      ...
      speech_mix_lengths: (Batch,), default None for chunk interator,
                      because the chunk-iterator does not have the
                      speech_lengths returned. see in
                      espnet2/iterators/chunk_iter_factory.py
      enroll_ref1: (Batch, samples_aux)
                          enrollment (raw audio or embedding) for speaker 1
      enroll_ref2: (Batch, samples_aux)
                          enrollment (raw audio or embedding) for speaker 2
      ...
      kwargs: "utt_id" is among the input.
docstr_tok: ["Frontend","+","Encoder","+","Decoder","+","Calc","loss","Args",":","speech_mix",":","(","Batch",",","samples",")","or","(","Batch",",","samples",",","channels",")","speech_ref1",":","(","Batch",",","samples",")","or","(","Batch",",","samples",",","channels",")","speech_ref2",":","(","Batch",",","samples",")","or","(","Batch",",","samples",",","channels",")","...","speech_mix_lengths",":","(","Batch",",",")",",","default","None","for","chunk","interator",",","because","the","chunk-iterator","does","not","have","the","speech_lengths","returned",".","see","in","espnet2\/iterators\/chunk_iter_factory.py","enroll_ref1",":","(","Batch",",","samples_aux",")","enrollment","(","raw","audio","or","embedding",")","for","speaker","1","enroll_ref2",":","(","Batch",",","samples_aux",")","enrollment","(","raw","audio","or","embedding",")","for","speaker","2","...","kwargs",":","``","utt_id","''","is","among","the","input","."]
code: def forward(
      self,
      speech_mix: torch.Tensor,
      speech_mix_lengths: torch.Tensor = None,
      **kwargs,
  ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
      """Frontend + Encoder + Decoder + Calc loss
  
      Args:
          speech_mix: (Batch, samples) or (Batch, samples, channels)
          speech_ref1: (Batch, samples)
                      or (Batch, samples, channels)
          speech_ref2: (Batch, samples)
                      or (Batch, samples, channels)
          ...
          speech_mix_lengths: (Batch,), default None for chunk interator,
                          because the chunk-iterator does not have the
                          speech_lengths returned. see in
                          espnet2/iterators/chunk_iter_factory.py
          enroll_ref1: (Batch, samples_aux)
                              enrollment (raw audio or embedding) for speaker 1
          enroll_ref2: (Batch, samples_aux)
                              enrollment (raw audio or embedding) for speaker 2
          ...
          kwargs: "utt_id" is among the input.
      """
      # reference speech signal of each speaker
      assert (
          "speech_ref1" in kwargs
      ), "At least 1 reference signal input is required."
      speech_ref = [
          kwargs.get(
              f"speech_ref{spk + 1}",
              torch.zeros_like(kwargs["speech_ref1"]),
          )
          for spk in range(self.num_spk)
          if "speech_ref{}".format(spk + 1) in kwargs
      ]
      num_spk = (
          len(speech_ref) if self.flexible_numspk else self.num_spk
      )
      assert len(speech_ref) == num_spk, (len(speech_ref), num_spk)
      # (Batch, num_speaker, samples) or (Batch, num_speaker, samples, channels)
      speech_ref = torch.stack(speech_ref, dim=1)
      batch_size = speech_mix.shape[0]
  
      assert (
          "enroll_ref1" in kwargs
      ), "At least 1 enrollment signal is required."
      # enrollment signal for each speaker (as the target)
      enroll_ref = [
          # (Batch, samples_aux)
          kwargs["enroll_ref{}".format(spk + 1)]
          for spk in range(num_spk)
          if "enroll_ref{}".format(spk + 1) in kwargs
      ]
      enroll_ref_lengths = [
          # (Batch,)
          kwargs.get(
              "enroll_ref{}_lengths".format(spk + 1),
              torch.ones(batch_size)
              .int()
              .fill_(enroll_ref[spk].size(1)),
          )
          for spk in range(num_spk)
          if "enroll_ref{}".format(spk + 1) in kwargs
      ]
  
      speech_lengths = (
          speech_mix_lengths
          if speech_mix_lengths is not None
          else torch.ones(batch_size).int().fill_(speech_mix.shape[1])
      )
      assert speech_lengths.dim() == 1, speech_lengths.shape
      # Check that batch_size is unified
      assert (
          speech_mix.shape[0]
          == speech_ref.shape[0]
          == speech_lengths.shape[0]
      ), (
          speech_mix.shape,
          speech_ref.shape,
          speech_lengths.shape,
      )
      for aux in enroll_ref:
          assert aux.shape[0] == speech_mix.shape[0], (
              aux.shape,
              speech_mix.shape,
          )
  
      # for data-parallel
      speech_ref = speech_ref[..., : speech_lengths.max()].unbind(dim=1)
  
      speech_mix = speech_mix[:, : speech_lengths.max()]
      enroll_ref = [
          enroll_ref[spk][:, : enroll_ref_lengths[spk].max()]
          for spk in range(len(enroll_ref))
      ]
      assert len(speech_ref) == len(enroll_ref), (
          len(speech_ref),
          len(enroll_ref),
      )
  
      additional = {}
      # Additional data for training the TSE model
      if self.flexible_numspk:
          additional["num_spk"] = num_spk
  
      # model forward
      (
          speech_pre,
          feature_mix,
          feature_pre,
          others,
      ) = self.forward_enhance(
          speech_mix,
          speech_lengths,
          enroll_ref,
          enroll_ref_lengths,
          additional,
      )
  
      # loss computation
      loss, stats, weight, perm = self.forward_loss(
          speech_pre,
          speech_lengths,
          feature_mix,
          feature_pre,
          others,
          speech_ref,
      )
      return loss, stats, weight
  
imports: import contextlib
  from typing import Dict, List, Optional, OrderedDict, Tuple
  import torch
  from typeguard import check_argument_types
  from espnet2.enh.decoder.abs_decoder import AbsDecoder
  from espnet2.enh.encoder.abs_encoder import AbsEncoder
  from espnet2.enh.extractor.abs_extractor import AbsExtractor
  from espnet2.enh.loss.criterions.tf_domain import FrequencyDomainLoss
  from espnet2.enh.loss.criterions.time_domain import TimeDomainLoss
  from espnet2.enh.loss.wrappers.abs_wrapper import AbsLossWrapper
  from espnet2.torch_utils.device_funcs import force_gatherable
  from espnet2.train.abs_espnet_model import AbsESPnetModel
label_desc: Use image node_id 2 for calling the ESPnetExtractionModel obj's underlying member method code with example usage: obj.forward(speech_mix, speech_mix_lengths) and returns: loss,  stats,  weight