class_docstr: ESPnet speech-to-speech translation model
code: def forward(
      self,
      src_speech: torch.Tensor,
      src_speech_lengths: torch.Tensor,
      tgt_speech: torch.Tensor,
      tgt_speech_lengths: torch.Tensor,
      tgt_text: Optional[torch.Tensor] = None,
      tgt_text_lengths: Optional[torch.Tensor] = None,
      src_text: Optional[torch.Tensor] = None,
      src_text_lengths: Optional[torch.Tensor] = None,
      spembs: Optional[torch.Tensor] = None,  # TODO(Jiatong)
      sids: Optional[torch.Tensor] = None,  # TODO(Jiatong)
      lids: Optional[torch.Tensor] = None,  # TODO(Jiatong)
      **kwargs,
  ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
      # TODO(jiatong): add comments etc.
  
      assert (
          src_speech.shape[0]
          == src_speech_lengths.shape[0]
          == tgt_speech.shape[0]
          == tgt_speech_lengths.shape[0]
      ), (
          src_speech.shape,
          src_speech_lengths.shape,
          tgt_speech.shape,
          tgt_speech_lengths.shape,
      )
  
      # additional checks with valid tgt_text and src_text
      if tgt_text is not None:
          assert tgt_text_lengths.dim() == 1, tgt_text_lengths.shape
          assert (
              src_speech.shape[0]
              == src_text.shape[0]
              == src_text_lengths.shape[0]
              == tgt_text.shape[0]
              == tgt_text_lengths.shape[0]
          ), (
              src_speech.shape,
              src_text.shape,
              src_text_lengths.shape,
              tgt_text.shape,
              tgt_text_lengths.shape,
          )
  
      batch_size = src_speech.shape[0]
      # for data-parallel
      src_speech = src_speech[:, : src_speech_lengths.max()]
      if src_text is not None:
          src_text = src_text[:, : src_text_lengths.max()]
      if tgt_text is not None:
          tgt_text = tgt_text[:, : tgt_text_lengths.max()]
      tgt_speech = tgt_speech[:, : tgt_speech_lengths.max()]
  
      # 0. Target feature extract
      # NOTE(jiatong): only for teaching-forcing in spectrogram
      if self.tgt_feats_extract is not None:
          tgt_feats, tgt_feats_lengths = self._extract_feats(
              tgt_speech, tgt_speech_lengths, target=True
          )
          # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
          if self.tgt_normalize is not None:
              tgt_feats, tgt_feats_lengths = self.tgt_normalize(
                  tgt_feats, tgt_feats_lengths
              )
      else:
          # NOTE(jiatong): for discrete unit case
          tgt_feats, tgt_feats_lengths = tgt_speech, tgt_speech_lengths
  
      # 1. Encoder
      if self.s2st_type == "discrete_unit":
          (
              encoder_out,
              inter_encoder_out,
          ), encoder_out_lens = self.encode(
              src_speech, src_speech_lengths, return_all_hs=True
          )
      else:
          encoder_out, encoder_out_lens = self.encode(
              src_speech, src_speech_lengths
          )
  
      loss_record = []
  
      ########################
      # Translaotron Forward #
      ########################
      if self.s2st_type == "translatotron":
          # use a shared encoder with three decoders (i.e., asr, st, s2st)
          # reference https://arxiv.org/pdf/1904.06037.pdf
  
          # asr_ctc
          if self.asr_ctc is not None and "asr_ctc" in self.losses:
              asr_ctc_loss, cer_asr_ctc = self._calc_ctc_loss(
                  encoder_out,
                  encoder_out_lens,
                  src_text,
                  src_text_lengths,
                  ctc_type="asr",
              )
              loss_record.append(
                  asr_ctc_loss * self.losses["asr_ctc"].weight
              )
          else:
              asr_ctc_loss, cer_asr_ctc = None, None
  
          # asr decoder
          if self.asr_decoder is not None and "src_attn" in self.losses:
              (
                  src_attn_loss,
                  acc_src_attn,
                  cer_src_attn,
                  wer_src_attn,
              ) = self._calc_asr_att_loss(
                  encoder_out,
                  encoder_out_lens,
                  src_text,
                  src_text_lengths,
              )
              loss_record.append(
                  src_attn_loss * self.losses["src_attn"].weight
              )
          else:
              (
                  src_attn_loss,
                  acc_src_attn,
                  cer_src_attn,
                  wer_src_attn,
              ) = (
                  None,
                  None,
                  None,
                  None,
              )
  
          # st decoder
          if self.st_decoder is not None and "tgt_attn" in self.losses:
              (
                  tgt_attn_loss,
                  acc_tgt_attn,
                  bleu_tgt_attn,
              ) = self._calc_st_att_loss(
                  encoder_out,
                  encoder_out_lens,
                  tgt_text,
                  tgt_text_lengths,
              )
              loss_record.append(
                  tgt_attn_loss * self.losses["tgt_attn"].weight
              )
          else:
              tgt_attn_loss, acc_tgt_attn, bleu_tgt_attn = (
                  None,
                  None,
                  None,
              )
  
          # NOTE(jiatong): the tgt_feats is also updated based on the reduction_factor
          (
              after_outs,
              before_outs,
              logits,
              att_ws,
              updated_tgt_feats,
              stop_labels,
              updated_tgt_feats_lengths,
          ) = self.synthesizer(
              encoder_out,
              encoder_out_lens,
              tgt_feats,
              tgt_feats_lengths,
              spembs,
              sids,
              lids,
          )
  
          syn_loss, l1_loss, mse_loss, bce_loss = self.losses[
              "synthesis"
          ](
              after_outs,
              before_outs,
              logits,
              updated_tgt_feats,
              stop_labels,
              updated_tgt_feats_lengths,
          )
          loss_record.append(syn_loss * self.losses["synthesis"].weight)
  
          # NOTE(jiatong): guided attention will be not used in multi-head attention
          if (
              "syn_guided_attn" in self.losses
              and self.synthesizer.atype != "multihead"
          ):
              # NOTE(kan-bayashi): length of output for auto-regressive
              # input will be changed when r > 1
              if self.synthesizer.reduction_factor > 1:
                  updated_tgt_feats_lengths_in = (
                      updated_tgt_feats_lengths.new(
                          [
                              olen // self.reduction_factor
                              for olen in updated_tgt_feats_lengths
                          ]
                      )
                  )
              else:
                  updated_tgt_feats_lengths_in = (
                      updated_tgt_feats_lengths
                  )
              syn_guided_attn_loss = self.losses["syn_guided_attn"](
                  att_ws=att_ws,
                  ilens=encoder_out_lens,
                  olens_in=updated_tgt_feats_lengths_in,
              )
              loss_record.append(
                  syn_guided_attn_loss
                  * self.losses["syn_guided_attn"].weight
              )
          else:
              syn_guided_attn_loss = None
  
          loss = sum(loss_record)
  
          stats = dict(
              loss=loss.item(),
              asr_ctc_loss=asr_ctc_loss.item()
              if asr_ctc_loss is not None
              else None,
              cer_asr_ctc=cer_asr_ctc,
              src_attn_loss=src_attn_loss.item()
              if src_attn_loss is not None
              else None,
              acc_src_attn=acc_src_attn,
              cer_src_attn=cer_src_attn,
              wer_src_attn=wer_src_attn,
              tgt_attn_loss=tgt_attn_loss.item()
              if tgt_attn_loss is not None
              else None,
              acc_tgt_attn=acc_tgt_attn,
              bleu_tgt_attn=bleu_tgt_attn,
              syn_loss=syn_loss.item()
              if syn_loss is not None
              else None,
              syn_guided_attn_loss=syn_guided_attn_loss.item()
              if syn_guided_attn_loss is not None
              else None,
              syn_l1_loss=l1_loss.item(),
              syn_mse_loss=mse_loss.item(),
              syn_bce_loss=bce_loss.item(),
          )
  
      #########################
      # Translaotron2 Forward #
      #########################
      elif self.s2st_type == "translatotron2":
          # use a sinlge decoder for synthesis
          # reference https://arxiv.org/pdf/2107.08661v5.pdf
  
          # asr_ctc
          if self.asr_ctc is not None and "asr_ctc" in self.losses:
              asr_ctc_loss, cer_asr_ctc = self._calc_ctc_loss(
                  encoder_out,
                  encoder_out_lens,
                  src_text,
                  src_text_lengths,
                  ctc_type="asr",
              )
              loss_record.append(
                  asr_ctc_loss * self.losses["asr_ctc"].weight
              )
          else:
              asr_ctc_loss, cer_asr_ctc = None, None
  
          # st decoder
          if self.st_decoder is not None and "tgt_attn" in self.losses:
              (
                  tgt_attn_loss,
                  acc_tgt_attn,
                  bleu_tgt_attn,
                  decoder_out,
                  _,
              ) = self._calc_st_att_loss(
                  encoder_out,
                  encoder_out_lens,
                  tgt_text,
                  tgt_text_lengths,
                  return_hs=True,
              )
              loss_record.append(
                  tgt_attn_loss * self.losses["tgt_attn"].weight
              )
          else:
              (
                  tgt_attn_loss,
                  acc_tgt_attn,
                  bleu_tgt_attn,
                  decoder_out,
              ) = (
                  None,
                  None,
                  None,
                  None,
              )
  
          assert (
              self.aux_attention is not None
          ), "must have aux attention in translatotron loss"
  
          # NOTE(jiatong): tgt_text_lengths + 1 for <eos>
          encoder_out_mask = (
              make_pad_mask(encoder_out_lens)
              .to(encoder_out.device)
              .unsqueeze(1)
          )
          attention_out = self.aux_attention(
              decoder_out,
              encoder_out,
              encoder_out,
              mask=encoder_out_mask,
          )
          decoder_out = torch.cat((decoder_out, attention_out), dim=-1)
  
          # NOTE(jiatong): the tgt_feats is also updated based on the reduction_factor
          # TODO(jiatong): use non-attentive tacotron-based synthesizer
          (
              after_outs,
              before_outs,
              logits,
              att_ws,
              updated_tgt_feats,
              stop_labels,
              updated_tgt_feats_lengths,
          ) = self.synthesizer(
              decoder_out,
              tgt_text_lengths + 1,  # NOTE(jiatong): +1 for <eos>
              tgt_feats,
              tgt_feats_lengths,
              spembs,
              sids,
              lids,
          )
  
          syn_loss, l1_loss, mse_loss, bce_loss = self.losses[
              "synthesis"
          ](
              after_outs,
              before_outs,
              logits,
              updated_tgt_feats,
              stop_labels,
              updated_tgt_feats_lengths,
          )
          # loss_record.append(syn_loss * self.losses["synthesis"].weight)
  
          loss = sum(loss_record)
  
          stats = dict(
              loss=loss.item(),
              asr_ctc_loss=asr_ctc_loss.item()
              if asr_ctc_loss is not None
              else None,
              cer_asr_ctc=cer_asr_ctc,
              tgt_attn_loss=tgt_attn_loss.item()
              if tgt_attn_loss is not None
              else None,
              acc_tgt_attn=acc_tgt_attn,
              bleu_tgt_attn=bleu_tgt_attn,
              syn_loss=syn_loss.item()
              if syn_loss is not None
              else None,
              syn_l1_loss=l1_loss.item()
              if l1_loss is not None
              else None,
              syn_mse_loss=mse_loss.item()
              if mse_loss is not None
              else None,
              syn_bce_loss=bce_loss.item()
              if bce_loss is not None
              else None,
          )
  
      #########################
      # Discrete unit Forward #
      #########################
      elif self.s2st_type == "discrete_unit":
          # discrete unit-based synthesis
          # Reference: https://arxiv.org/pdf/2107.05604.pdf
  
          encoder_layer_for_asr = len(inter_encoder_out) // 2
          encoder_layer_for_st = len(inter_encoder_out) * 2 // 3
  
          # asr_ctc
          if self.asr_ctc is not None and "asr_ctc" in self.losses:
              asr_ctc_loss, cer_asr_ctc = self._calc_ctc_loss(
                  inter_encoder_out[encoder_layer_for_asr],
                  encoder_out_lens,
                  src_text,
                  src_text_lengths,
                  ctc_type="asr",
              )
              loss_record.append(
                  asr_ctc_loss * self.losses["asr_ctc"].weight
              )
          else:
              asr_ctc_loss, cer_asr_ctc = None, None
  
          # asr decoder
          if self.asr_decoder is not None and "src_attn" in self.losses:
              (
                  src_attn_loss,
                  acc_src_attn,
                  cer_src_attn,
                  wer_src_attn,
              ) = self._calc_asr_att_loss(
                  encoder_out,
                  encoder_out_lens,
                  src_text,
                  src_text_lengths,
              )
              loss_record.append(
                  src_attn_loss * self.losses["src_attn"].weight
              )
          else:
              (
                  src_attn_loss,
                  acc_src_attn,
                  cer_src_attn,
                  wer_src_attn,
              ) = (
                  None,
                  None,
                  None,
                  None,
              )
  
          # st decoder
          if self.st_decoder is not None and "tgt_attn" in self.losses:
              (
                  tgt_attn_loss,
                  acc_tgt_attn,
                  bleu_tgt_attn,
              ) = self._calc_st_att_loss(
                  inter_encoder_out[encoder_layer_for_st],
                  encoder_out_lens,
                  tgt_text,
                  tgt_text_lengths,
              )
              loss_record.append(
                  tgt_attn_loss * self.losses["tgt_attn"].weight
              )
          else:
              (
                  tgt_attn_loss,
                  acc_tgt_attn,
                  bleu_tgt_attn,
                  decoder_out,
              ) = (
                  None,
                  None,
                  None,
                  None,
              )
  
          # synthesizer
          (
              unit_attn_loss,
              acc_unit_attn,
              syn_hidden,
              syn_hidden_lengths,
          ) = self._calc_unit_att_loss(
              encoder_out,
              encoder_out_lens,
              tgt_speech,
              tgt_speech_lengths,
              return_all_hs=True,
          )
  
          loss_record.append(
              unit_attn_loss * self.losses["synthesis"].weight
          )
  
          unit_decoder_layer_for_st = len(syn_hidden) // 2
  
          # st_ctc
          if self.st_ctc is not None and "st_ctc" in self.losses:
              st_ctc_loss, cer_st_ctc = self._calc_ctc_loss(
                  syn_hidden[unit_decoder_layer_for_st],
                  tgt_speech_lengths + 1,
                  tgt_text,
                  tgt_text_lengths,
                  ctc_type="st",
              )
              loss_record.append(
                  st_ctc_loss * self.losses["st_ctc"].weight
              )
          else:
              st_ctc_loss, cer_st_ctc = None, None
  
          loss = sum(loss_record)
  
          stats = dict(
              loss=loss.item(),
              asr_ctc_loss=asr_ctc_loss.item()
              if asr_ctc_loss is not None
              else None,
              cer_asr_ctc=cer_asr_ctc,
              src_attn_loss=src_attn_loss.item()
              if src_attn_loss is not None
              else None,
              acc_src_attn=acc_src_attn,
              cer_src_attn=cer_src_attn,
              wer_src_attn=wer_src_attn,
              tgt_attn_loss=tgt_attn_loss.item()
              if tgt_attn_loss is not None
              else None,
              acc_tgt_attn=acc_tgt_attn,
              bleu_tgt_attn=bleu_tgt_attn,
              st_ctc_loss=st_ctc_loss.item()
              if st_ctc_loss is not None
              else None,
              cer_st_ctc=cer_st_ctc,
              unit_attn_loss=unit_attn_loss.item()
              if unit_attn_loss is not None
              else None,
              acc_unit_attn=acc_unit_attn
              if acc_unit_attn is not None
              else None,
          )
  
      #################
      # Unity Forward #
      #################
      elif self.s2st_type == "unity":
          # unity
          # Reference: https://arxiv.org/pdf/2212.08055.pdf
  
          # asr_ctc
          if self.asr_ctc is not None and "asr_ctc" in self.losses:
              asr_ctc_loss, cer_asr_ctc = self._calc_ctc_loss(
                  encoder_out,
                  encoder_out_lens,
                  src_text,
                  src_text_lengths,
                  ctc_type="asr",
              )
              loss_record.append(
                  asr_ctc_loss * self.losses["asr_ctc"].weight
              )
          else:
              asr_ctc_loss, cer_asr_ctc = None, None
  
          # st decoder
          assert (
              self.st_decoder is not None and "tgt_attn" in self.losses
          ), "st_decoder is necessary for unity-based model"
          (
              tgt_attn_loss,
              acc_tgt_attn,
              bleu_tgt_attn,
              decoder_out,
              _,
          ) = self._calc_st_att_loss(
              encoder_out,
              encoder_out_lens,
              tgt_text,
              tgt_text_lengths,
              return_hs=True,
          )
          loss_record.append(
              tgt_attn_loss * self.losses["tgt_attn"].weight
          )
  
          assert (
              self.unit_encoder is not None
          ), "unit_encoder is necessary for unity-based model"
  
          (
              unit_encoder_out,
              unit_encoder_out_lengths,
              _,
          ) = self.unit_encoder(decoder_out, tgt_text_lengths + 1)
  
          # synthesizer
          unit_attn_loss, acc_unit_attn = self._calc_unit_att_loss(
              unit_encoder_out,
              unit_encoder_out_lengths,
              tgt_speech,
              tgt_speech_lengths,
          )
          loss_record.append(
              unit_attn_loss * self.losses["synthesis"].weight
          )
  
          loss = sum(loss_record)
  
          stats = dict(
              loss=loss.item(),
              asr_ctc_loss=asr_ctc_loss.item()
              if asr_ctc_loss is not None
              else None,
              cer_asr_ctc=cer_asr_ctc,
              tgt_attn_loss=tgt_attn_loss.item()
              if tgt_attn_loss is not None
              else None,
              acc_tgt_attn=acc_tgt_attn,
              bleu_tgt_attn=bleu_tgt_attn,
              unit_attn_loss=unit_attn_loss.item()
              if unit_attn_loss is not None
              else None,
              acc_unit_attn=acc_unit_attn
              if acc_unit_attn is not None
              else None,
          )
  
      else:
          raise ValueError("Not supported s2st type {}")
  
      # force_gatherable: to-device and to-tensor if scalar for DataParallel
      loss, stats, weight = force_gatherable(
          (loss, stats, batch_size), loss.device
      )
      return loss, stats, weight
  
imports: import logging
  from contextlib import contextmanager
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from packaging.version import parse
  from typeguard import check_argument_types
  from espnet2.asr.ctc import CTC
  from espnet2.asr.decoder.abs_decoder import AbsDecoder
  from espnet2.asr.encoder.abs_encoder import AbsEncoder
  from espnet2.asr.encoder.conformer_encoder import ConformerEncoder
  from espnet2.asr.encoder.transformer_encoder import TransformerEncoder
  from espnet2.asr.frontend.abs_frontend import AbsFrontend
  from espnet2.asr.postencoder.abs_postencoder import AbsPostEncoder
  from espnet2.asr.preencoder.abs_preencoder import AbsPreEncoder
  from espnet2.asr.specaug.abs_specaug import AbsSpecAug
  from espnet2.layers.abs_normalize import AbsNormalize
  from espnet2.s2st.aux_attention.abs_aux_attention import AbsS2STAuxAttention
  from espnet2.s2st.losses.abs_loss import AbsS2STLoss
  from espnet2.s2st.synthesizer.abs_synthesizer import AbsSynthesizer
  from espnet2.s2st.tgt_feats_extract.abs_tgt_feats_extract import AbsTgtFeatsExtract
  from espnet2.torch_utils.device_funcs import force_gatherable
  from espnet2.train.abs_espnet_model import AbsESPnetModel
  from espnet.nets.e2e_asr_common import ErrorCalculator
  from espnet.nets.e2e_mt_common import ErrorCalculator
  from espnet.nets.pytorch_backend.nets_utils import make_pad_mask, th_accuracy
  from espnet.nets.pytorch_backend.transformer.add_sos_eos import add_sos_eos
label_desc: Use image node_id 2 for calling the ESPnetS2STModel obj's underlying member method code with example usage: obj.forward(src_speech, src_speech_lengths, tgt_speech, tgt_speech_lengths, tgt_text, tgt_text_lengths, src_text, src_text_lengths, spembs, sids, lids) and returns: loss,  stats,  weight