text: """Auxiliary task implementation for transducer models."""
  
  from itertools import chain
  from typing import List
  from typing import Tuple
  from typing import Union
  
  import torch
  import torch.nn.functional as F
  
  from espnet.nets.transducer_decoder_interface import TransducerDecoderInterface
  
  
  class AuxiliaryTask(torch.nn.Module):
      """Auxiliary task module."""
  
      def __init__(
          self,
          decoder: Union[torch.nn.Module, TransducerDecoderInterface],
          joint_network: torch.nn.Module,
          rnnt_criterion: torch.nn.Module,
          aux_task_type: str,
          aux_task_weight: int,
          encoder_out_dim: int,
          joint_dim: int,
      ):
          """Auxiliary task initialization.
  
          Args:
              decoder: Decoder module
              joint_network: Joint network module
              aux_task_type: Auxiliary task type
              aux_task_weight: Auxiliary task weight
              encoder_out: Encoder output dimension
              joint_dim: Joint space dimension
  
          """
          super().__init__()
  
          self.rnnt_criterion = rnnt_criterion
  
          self.mlp_net = torch.nn.Sequential(
              torch.nn.Linear(encoder_out_dim, joint_dim),
              torch.nn.ReLU(),
              torch.nn.Linear(joint_dim, joint_dim),
          )
  
          self.decoder = decoder
          self.joint_network = joint_network
  
          self.aux_task_type = aux_task_type
          self.aux_task_weight = aux_task_weight
  
      def forward(
          self,
          enc_out_aux: List,
          dec_out: torch.Tensor,
          main_joint: torch.Tensor,
          target: torch.Tensor,
          pred_len: torch.Tensor,
          target_len: torch.Tensor,
      ) -> Tuple[torch.Tensor, torch.Tensor]:
          """Forward auxiliary task.
  
          Args:
              enc_out_aux: List of encoder intermediate outputs
              dec_out: Decoder outputs
              main_joint: Joint output for main task
              target: Target labels
              pred_len: Prediction lengths
              target_len: Target lengths
  
          Returns:
              : (Weighted auxiliary transducer loss, Weighted auxiliary symmetric KL loss)
  
          """
          aux_trans = 0
          aux_symm_kl = 0
  
          for p in chain(self.decoder.parameters(), self.joint_network.parameters()):
              p.requires_grad = False
  
          for i, enc_aux in enumerate(enc_out_aux):
              aux_mlp = self.mlp_net(enc_aux)
  
              aux_joint = self.joint_network(
                  aux_mlp.unsqueeze(2),
                  dec_out.unsqueeze(1),
                  is_aux=True,
              )
  
              if self.aux_task_type != "symm_kl_div":
                  aux_trans += self.rnnt_criterion(
                      aux_joint,
                      target,
                      pred_len,
                      target_len,
                  )
  
              if self.aux_task_type != "default":
                  aux_symm_kl += F.kl_div(
                      F.log_softmax(main_joint, dim=-1),
                      F.softmax(aux_joint, dim=-1),
                      reduction="mean",
                  ) + F.kl_div(
                      F.log_softmax(aux_joint, dim=-1),
                      F.softmax(main_joint, dim=-1),
                      reduction="mean",
                  )
  
          for p in chain(self.decoder.parameters(), self.joint_network.parameters()):
              p.requires_grad = True
  
          return self.aux_task_weight * aux_trans, self.aux_task_weight * aux_symm_kl
  