docstr: Extend `_forward` for (possibly) multiple decoder pass:
  Always a default (future masked) decoder forward pass,
  Possibly a second future aware decoder pass for joint learn
  full context alignement, :cite:`garg2019jointly`.
  
  Args:
      * All arguments of _forward, of which
      with_align (bool): needed to compute attn_align
      return_attn (bool): to force MHA to return attns
  
  Returns:
      (FloatTensor, FloatTensor, FloatTensor or None):
  
      * layer_out ``(batch_size, T, model_dim)``
      * top_attn ``(batch_size, T, src_len)``
      * attn_align ``(batch_size, T, src_len)`` or None
docstr_tok: ["Extend","`","_forward","`","for","(","possibly",")","multiple","decoder","pass",":","Always","a","default","(","future","masked",")","decoder","forward","pass",",","Possibly","a","second","future","aware","decoder","pass","for","joint","learn","full","context","alignement",",",":","cite",":","`","garg2019jointly","`",".","Args",":","*","All","arguments","of","_forward",",","of","which","with_align","(","bool",")",":","needed","to","compute","attn_align","return_attn","(","bool",")",":","to","force","MHA","to","return","attns","Returns",":","(","FloatTensor",",","FloatTensor",",","FloatTensor","or","None",")",":","*","layer_out","``","(","batch_size",",","T",",","model_dim",")","``","*","top_attn","``","(","batch_size",",","T",",","src_len",")","``","*","attn_align","``","(","batch_size",",","T",",","src_len",")","``","or","None"]
code: def forward(self, *args, **kwargs):
      """Extend `_forward` for (possibly) multiple decoder pass:
      Always a default (future masked) decoder forward pass,
      Possibly a second future aware decoder pass for joint learn
      full context alignement, :cite:`garg2019jointly`.
  
      Args:
          * All arguments of _forward, of which
          with_align (bool): needed to compute attn_align
          return_attn (bool): to force MHA to return attns
  
      Returns:
          (FloatTensor, FloatTensor, FloatTensor or None):
  
          * layer_out ``(batch_size, T, model_dim)``
          * top_attn ``(batch_size, T, src_len)``
          * attn_align ``(batch_size, T, src_len)`` or None
      """
      with_align = kwargs.pop("with_align", False)
      layer_out, attns = self._forward(*args, **kwargs)
      top_attn = (
          None if attns is None else attns[:, 0, :, :].contiguous()
      )
      attn_align = None
      if with_align:
          if self.full_context_alignment:
              # return _, (B, Q_len, K_len)
              _, attns = self._forward(*args, **kwargs, future=True)
  
          if self.alignment_heads > 0:
              attns = attns[
                  :, : self.alignment_heads, :, :
              ].contiguous()
          # layer average attention across heads, get ``(B, Q, K)``
          # Case 1: no full_context, no align heads -> layer avg baseline
          # Case 2: no full_context, 1 align heads -> guided align
          # Case 3: full_context, 1 align heads -> full cte guided align
          attn_align = attns.mean(dim=1)
      return layer_out, top_attn, attn_align
  
imports: import torch
  import torch.nn
  from onmt.decoders.decoder import DecoderBase
  from onmt.modules import MultiHeadedAttention, AverageAttention
  from onmt.modules.position_ffn import PositionwiseFeedForward
  from onmt.modules.position_ffn import ActivationFunction
  from onmt.modules.moe import MoE
  from onmt.utils.misc import sequence_mask
  from onmt.modules.rmsnorm import RMSNorm
label_desc: Use image node_id 2 for calling the TransformerDecoderLayerBase obj's underlying member method code with example usage: obj.forward() and returns: layer_out,  top_attn,  attn_align