class_docstr: MEGA decoder module.
  
  Based on https://arxiv.org/pdf/2209.10655.pdf.
  
  Args:
      vocab_size: Vocabulary size.
      block_size: Input/Output size.
      linear_size: NormalizedPositionwiseFeedForward hidden size.
      qk_size: Shared query and key size for attention module.
      v_size: Value size for attention module.
      num_heads: Number of EMA heads.
      rel_pos_bias: Type of relative position bias in attention module.
      max_positions: Maximum number of position for RelativePositionBias.
      truncation_length: Maximum length for truncation in EMA module.
      normalization_type: Normalization layer type.
      normalization_args: Normalization layer arguments.
      activation_type: Activation function type.
      activation_args: Activation function arguments.
      chunk_size: Chunk size for attention computation (-1 = full context).
      num_blocks: Number of MEGA blocks.
      dropout_rate: Dropout rate for MEGA internal modules.
      embed_dropout_rate: Dropout rate for embedding layer.
      att_dropout_rate: Dropout rate for the attention module.
      ema_dropout_rate: Dropout rate for the EMA module.
      ffn_dropout_rate: Dropout rate for the feed-forward module.
      embed_pad: Embedding padding symbol ID.
docstr: Encode source label sequences.
  
  Args:
      labels: Decoder input sequences. (B, L)
  
  Returns:
      out: Decoder output sequences. (B, U, D_dec)
docstr_tok: ["Encode","source","label","sequences",".","Args",":","labels",":","Decoder","input","sequences",".","(","B",",","L",")","Returns",":","out",":","Decoder","output","sequences",".","(","B",",","U",",","D_dec",")"]
code: def forward(self, labels: torch.Tensor) -> torch.Tensor:
      """Encode source label sequences.
  
      Args:
          labels: Decoder input sequences. (B, L)
  
      Returns:
          out: Decoder output sequences. (B, U, D_dec)
  
      """
      batch, length = labels.size()
  
      if 0 < self.chunk_size < length and length % self.chunk_size != 0:
          num_paddings = (
              math.ceil(length / self.chunk_size) * self.chunk_size
              - length
          )
          labels = torch.nn.functional.pad(
              labels, (0, num_paddings), value=self.pad_idx
          )
      else:
          num_paddings = 0
  
      mask = (labels == self.pad_idx).unsqueeze(1)
      mask[..., 0] = False
      mask = mask.to(device=labels.device, dtype=torch.bool)
  
      _length = (
          self.chunk_size if 0 < self.chunk_size < length else length
      )
  
      attn_mask = torch.ones(
          (_length, _length), device=labels.device, dtype=torch.bool
      )
      attn_mask = torch.triu(attn_mask, 1, out=attn_mask).unsqueeze(0)
  
      x = self.dropout_embed(self.embed(labels)).transpose(0, 1)
  
      for idx, (mega_block, nffn) in enumerate(self.mega_blocks):
          x, _ = mega_block(x, mask=mask, attn_mask=attn_mask)
  
          x = nffn(x)
  
      out = self.final_norm(x).transpose(0, 1)
  
      if num_paddings > 0:
          out = out[:, :length, :]
  
      return out
  
imports: import math
  from typing import Dict, List, Optional, Tuple
  import torch
  from typeguard import check_argument_types
  from espnet2.asr_transducer.activation import get_activation
  from espnet2.asr_transducer.beam_search_transducer import Hypothesis
  from espnet2.asr_transducer.decoder.abs_decoder import AbsDecoder
  from espnet2.asr_transducer.decoder.blocks.mega import MEGA
  from espnet2.asr_transducer.decoder.modules.mega.feed_forward import NormalizedPositionwiseFeedForward
  from espnet2.asr_transducer.normalization import get_normalization
label_desc: Use image node_id 2 for calling the MEGADecoder obj's underlying member method code with example usage: obj.forward(labels) and returns: out