class_docstr: Speaker Diarization model
  
  If "attractor" is "None", SA-EEND will be used.
  Else if "attractor" is not "None", EEND-EDA will be used.
  For the details about SA-EEND and EEND-EDA, refer to the following papers:
  SA-EEND: https://arxiv.org/pdf/1909.06247.pdf
  EEND-EDA: https://arxiv.org/pdf/2005.09921.pdf, https://arxiv.org/pdf/2106.10654.pdf
docstr: Frontend + Encoder + Decoder + Calc loss
  
  Args:
      speech: (Batch, samples)
      speech_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
      spk_labels: (Batch, )
      kwargs: "utt_id" is among the input.
docstr_tok: ["Frontend","+","Encoder","+","Decoder","+","Calc","loss","Args",":","speech",":","(","Batch",",","samples",")","speech_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","spk_labels",":","(","Batch",",",")","kwargs",":","``","utt_id","''","is","among","the","input","."]
code: def forward(
      self,
      speech: torch.Tensor,
      speech_lengths: torch.Tensor = None,
      spk_labels: torch.Tensor = None,
      spk_labels_lengths: torch.Tensor = None,
      **kwargs,
  ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
      """Frontend + Encoder + Decoder + Calc loss
  
      Args:
          speech: (Batch, samples)
          speech_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
          spk_labels: (Batch, )
          kwargs: "utt_id" is among the input.
      """
      assert speech.shape[0] == spk_labels.shape[0], (
          speech.shape,
          spk_labels.shape,
      )
      batch_size = speech.shape[0]
  
      # 1. Encoder
      # Use bottleneck_feats if exist. Only for "enh + diar" task.
      bottleneck_feats = kwargs.get("bottleneck_feats", None)
      bottleneck_feats_lengths = kwargs.get(
          "bottleneck_feats_lengths", None
      )
      encoder_out, encoder_out_lens = self.encode(
          speech,
          speech_lengths,
          bottleneck_feats,
          bottleneck_feats_lengths,
      )
  
      if self.attractor is None:
          # 2a. Decoder (baiscally a predction layer after encoder_out)
          pred = self.decoder(encoder_out, encoder_out_lens)
      else:
          # 2b. Encoder Decoder Attractors
          # Shuffle the chronological order of encoder_out, then calculate attractor
          encoder_out_shuffled = encoder_out.clone()
          for i in range(len(encoder_out_lens)):
              encoder_out_shuffled[
                  i, : encoder_out_lens[i], :
              ] = encoder_out[i, torch.randperm(encoder_out_lens[i]), :]
          attractor, att_prob = self.attractor(
              encoder_out_shuffled,
              encoder_out_lens,
              to_device(
                  self,
                  torch.zeros(
                      encoder_out.size(0),
                      spk_labels.size(2) + 1,
                      encoder_out.size(2),
                  ),
              ),
          )
          # Remove the final attractor which does not correspond to a speaker
          # Then multiply the attractors and encoder_out
          pred = torch.bmm(
              encoder_out, attractor[:, :-1, :].permute(0, 2, 1)
          )
      # 3. Aggregate time-domain labels
      spk_labels, spk_labels_lengths = self.label_aggregator(
          spk_labels, spk_labels_lengths
      )
  
      # If encoder uses conv* as input_layer (i.e., subsampling),
      # the sequence length of 'pred' might be slighly less than the
      # length of 'spk_labels'. Here we force them to be equal.
      length_diff_tolerance = 2
      length_diff = spk_labels.shape[1] - pred.shape[1]
      if length_diff > 0 and length_diff <= length_diff_tolerance:
          spk_labels = spk_labels[:, 0 : pred.shape[1], :]
  
      if self.attractor is None:
          loss_pit, loss_att = None, None
          loss, perm_idx, perm_list, label_perm = self.pit_loss(
              pred, spk_labels, encoder_out_lens
          )
      else:
          loss_pit, perm_idx, perm_list, label_perm = self.pit_loss(
              pred, spk_labels, encoder_out_lens
          )
          loss_att = self.attractor_loss(att_prob, spk_labels)
          loss = (
              self.diar_weight * loss_pit
              + self.attractor_weight * loss_att
          )
      (
          correct,
          num_frames,
          speech_scored,
          speech_miss,
          speech_falarm,
          speaker_scored,
          speaker_miss,
          speaker_falarm,
          speaker_error,
      ) = self.calc_diarization_error(
          pred, label_perm, encoder_out_lens
      )
  
      if speech_scored > 0 and num_frames > 0:
          sad_mr, sad_fr, mi, fa, cf, acc, der = (
              speech_miss / speech_scored,
              speech_falarm / speech_scored,
              speaker_miss / speaker_scored,
              speaker_falarm / speaker_scored,
              speaker_error / speaker_scored,
              correct / num_frames,
              (speaker_miss + speaker_falarm + speaker_error)
              / speaker_scored,
          )
      else:
          sad_mr, sad_fr, mi, fa, cf, acc, der = 0, 0, 0, 0, 0, 0, 0
  
      stats = dict(
          loss=loss.detach(),
          loss_att=loss_att.detach() if loss_att is not None else None,
          loss_pit=loss_pit.detach() if loss_pit is not None else None,
          sad_mr=sad_mr,
          sad_fr=sad_fr,
          mi=mi,
          fa=fa,
          cf=cf,
          acc=acc,
          der=der,
      )
  
      loss, stats, weight = force_gatherable(
          (loss, stats, batch_size), loss.device
      )
      return loss, stats, weight
  
imports: from contextlib import contextmanager
  from itertools import permutations
  from typing import Dict, Optional, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from packaging.version import parse
  from typeguard import check_argument_types
  from espnet2.asr.encoder.abs_encoder import AbsEncoder
  from espnet2.asr.frontend.abs_frontend import AbsFrontend
  from espnet2.asr.specaug.abs_specaug import AbsSpecAug
  from espnet2.diar.attractor.abs_attractor import AbsAttractor
  from espnet2.diar.decoder.abs_decoder import AbsDecoder
  from espnet2.layers.abs_normalize import AbsNormalize
  from espnet2.torch_utils.device_funcs import force_gatherable
  from espnet2.train.abs_espnet_model import AbsESPnetModel
  from espnet.nets.pytorch_backend.nets_utils import to_device
label_desc: Use image node_id 2 for calling the ESPnetDiarizationModel obj's underlying member method code with example usage: obj.forward(speech, speech_lengths, spk_labels, spk_labels_lengths) and returns: loss,  stats,  weight