class_docstr: Whole utterance augmented memory multihead attention using BMM.
  
  Different with previous augmented memory multihead attention where
  the utterance is chunked into segments. Here we use attention mask
  achieve so. The input embedding [right_context, utterance, summary]
  is a concatenation of right context, utterance and summary.
  
  Right context block is the concatenation of all the right context for
  each segments. [right_context_0, right_context_1, ..., right_context_n]
  For example, if we have utterance = [v0, v1, v2, ...., v20]. segment
  size 8, right_context size 4. Then the right context blocks =
  [v8, v9, v10, v11, v16, v17, v18, v19, 0, 0, 0, 0], where v8, v9, v10,
  and v11 are the right context for first segment. v16, v17, v18 and v19
  are the right context for second segment. 0, 0, 0 and 0 are right context
  for the last segment.
  
  utterance is corresponding to input embedding sequence
  
  summary is concatenation of average of each segments. [summary_0,
  summary_1, ..., ].
  
  In augmented memory multihead attention, the query is [right_context,
  utterance, summary], key is [memory, right_context, utterance]. Different
  with AugmentedMemoryMultiheadAttentionBmm, memory here is passed from
  previous attention layer. For the first attention layer, memory is average
  of each segment.
  
  Memory is a concatenation of memory from each segments in previous attention
  layer. For example, current layer is i, then memory is [m_0, m_1, ..., m_n].
  Each m_k is the output from seg_k in layer i-1.
  
  args:
      input_dim: input embedding dimension
      num_heads: number of heads in multihead self-attention
      dropout: attention dropout
      std_scale: if std_scale is not None. The weak attention suppression is
          turned on. For std_scale = 0.5, all the attention smaller than
          mean + 0.5 * std will be suppressed.
      scaled_init: whether to use scaled init for linear weight
      tanh_on_mem: whether to use tanh on memory output
      use_mem: whether to use memory or not. When max_memory_size is 0, then
          we don't have memory anymore.
      layer_index: current self-attention layer index that is used in depth
          initialization
      max_relative_position: max relative position used in relative position
          embedding
      rpe_old_option: To be compatible with previous model. The previous model
          was trained with attention += attention + rpe. The correct equation
          should be attention = attention + rpe
class_docstr_tok: ["Whole","utterance","augmented","memory","multihead","attention","using","BMM",".","Different","with","previous","augmented","memory","multihead","attention","where","the","utterance","is","chunked","into","segments",".","Here","we","use","attention","mask","achieve","so",".","The","input","embedding","[","right_context",",","utterance",",","summary","]","is","a","concatenation","of","right","context",",","utterance","and","summary",".","Right","context","block","is","the","concatenation","of","all","the","right","context","for","each","segments",".","[","right_context_0",",","right_context_1",",","...",",","right_context_n","]","For","example",",","if","we","have","utterance","=","[","v0",",","v1",",","v2",",","....",",","v20","]",".","segment","size","8",",","right_context","size","4",".","Then","the","right","context","blocks","=","[","v8",",","v9",",","v10",",","v11",",","v16",",","v17",",","v18",",","v19",",","0",",","0",",","0",",","0","]",",","where","v8",",","v9",",","v10",",","and","v11","are","the","right","context","for","first","segment",".","v16",",","v17",",","v18","and","v19","are","the","right","context","for","second","segment",".","0",",","0",",","0","and","0","are","right","context","for","the","last","segment",".","utterance","is","corresponding","to","input","embedding","sequence","summary","is","concatenation","of","average","of","each","segments",".","[","summary_0",",","summary_1",",","...",",","]",".","In","augmented","memory","multihead","attention",",","the","query","is","[","right_context",",","utterance",",","summary","]",",","key","is","[","memory",",","right_context",",","utterance","]",".","Different","with","AugmentedMemoryMultiheadAttentionBmm",",","memory","here","is","passed","from","previous","attention","layer",".","For","the","first","attention","layer",",","memory","is","average","of","each","segment",".","Memory","is","a","concatenation","of","memory","from","each","segments","in","previous","attention","layer",".","For","example",",","current","layer","is","i",",","then","memory","is","[","m_0",",","m_1",",","...",",","m_n","]",".","Each","m_k","is","the","output","from","seg_k","in","layer","i-1",".","args",":","input_dim",":","input","embedding","dimension","num_heads",":","number","of","heads","in","multihead","self-attention","dropout",":","attention","dropout","std_scale",":","if","std_scale","is","not","None",".","The","weak","attention","suppression","is","turned","on",".","For","std_scale","=","0.5",",","all","the","attention","smaller","than","mean","+","0.5","*","std","will","be","suppressed",".","scaled_init",":","whether","to","use","scaled","init","for","linear","weight","tanh_on_mem",":","whether","to","use","tanh","on","memory","output","use_mem",":","whether","to","use","memory","or","not",".","When","max_memory_size","is","0",",","then","we","do","n't","have","memory","anymore",".","layer_index",":","current","self-attention","layer","index","that","is","used","in","depth","initialization","max_relative_position",":","max","relative","position","used","in","relative","position","embedding","rpe_old_option",":","To","be","compatible","with","previous","model",".","The","previous","model","was","trained","with","attention","+=","attention","+","rpe",".","The","correct","equation","should","be","attention","=","attention","+","rpe"]
docstr: forward function for NoSegAugmentedMemoryMultiheadAttentionBmm in training.
  
  args:
      input: formed in the following way
          [right_context_0, right_contex_1, ..., seg_0, seg_1,
          ..., summary_0, summary_1,..]
      lengths: the length of query which is [seg_0, seg_1, ....]
      mems: [mem_0, mem_1, ...].
      attention_mask: attention mask for query = [right_context, query, summary]
          key = [mem, right_context, query]. This is only used for traing.
docstr_tok: ["forward","function","for","NoSegAugmentedMemoryMultiheadAttentionBmm","in","training",".","args",":","input",":","formed","in","the","following","way","[","right_context_0",",","right_contex_1",",","...",",","seg_0",",","seg_1",",","...",",","summary_0",",","summary_1",",","..","]","lengths",":","the","length","of","query","which","is","[","seg_0",",","seg_1",",","....","]","mems",":","[","mem_0",",","mem_1",",","...","]",".","attention_mask",":","attention","mask","for","query","=","[","right_context",",","query",",","summary","]","key","=","[","mem",",","right_context",",","query","]",".","This","is","only","used","for","traing","."]
code: def forward(
      self,
      input: Tensor,
      lengths: Tensor,
      mems: Tensor,
      attention_mask: Tensor,
      pre_mems: Optional[Tensor] = None,
      left_context_key: Optional[Tensor] = None,
      left_context_val: Optional[Tensor] = None,
      rpe: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
      """
      forward function for NoSegAugmentedMemoryMultiheadAttentionBmm in training.
  
      args:
          input: formed in the following way
              [right_context_0, right_contex_1, ..., seg_0, seg_1,
              ..., summary_0, summary_1,..]
          lengths: the length of query which is [seg_0, seg_1, ....]
          mems: [mem_0, mem_1, ...].
          attention_mask: attention mask for query = [right_context, query, summary]
              key = [mem, right_context, query]. This is only used for traing.
  
      """
      if self.use_mem:
          mem_length = mems.size(0)
          summary_length = mem_length + 1
          if pre_mems is not None:
              mems = torch.cat([pre_mems, mems], dim=0)
      else:
          mem_length = 0
          summary_length = 0
  
      # In training, lc_length = 0
      if left_context_key is not None:
          lc_length = left_context_key.size(0)
      else:
          lc_length = 0
      results = self.prepare_qkv(
          input=input,
          mems=mems,
          lengths=lengths,
          summary_length=summary_length,
          lc_length=lc_length,
      )
      (
          result_qkv,
          input_shape,
          result_lengths_info,
          padding_mask,
      ) = results
      q, k, v = result_qkv
      (
          mem_length,
          utterance_length,
          right_context_blocks_length,
          key_length,
      ) = result_lengths_info
  
      if left_context_key is not None:
          # add the cache key and value
          new_k = torch.cat(
              [
                  k[: mem_length + right_context_blocks_length, :, :],
                  left_context_key,
                  k[-utterance_length:, :, :],
              ],
              dim=0,
          )
          new_v = torch.cat(
              [
                  v[: mem_length + right_context_blocks_length, :, :],
                  left_context_val,
                  v[-utterance_length:, :, :],
              ],
              dim=0,
          )
          next_k = new_k[
              mem_length + right_context_blocks_length :, :, :
          ]
          next_v = new_v[
              mem_length + right_context_blocks_length :, :, :
          ]
      else:
          new_k = k
          new_v = v
          next_k = None
          next_v = None
  
      (
          attention_weights,
          attention_weights_float,
          v,
      ) = self.prepare_attention_weights(
          q=q,
          new_k=new_k,
          new_v=new_v,
          input_shape=input_shape,
          rpe=rpe,
      )
  
      # mask attention
      attention_mask = attention_mask.unsqueeze(0)
      attention_weights_float = attention_weights_float.masked_fill(
          attention_mask, float(self.negative_inf)
      )
  
      rc_output_memory = self.prepare_attention_output(
          attention_weights=attention_weights,
          attention_weights_float=attention_weights_float,
          v=v,
          input_shape=input_shape,
          key_length=key_length,
          padding_mask=padding_mask,
          rpe=rpe,
      )
  
      if self.use_mem:
          # next_m length equals to summary length - 1
          # last memory is ignored
          if self.mini_batches:
              next_m = rc_output_memory[-summary_length:]
          else:
              next_m = rc_output_memory[-summary_length:-1]
  
          next_m = self.squash_mem(next_m)
          # rc and output
          rc_output = rc_output_memory[:-summary_length]
          if not self.nonlinear_squash_mem:
              next_m = torch.clamp(next_m, min=-10, max=10)
      else:
          next_m = mems
          rc_output = rc_output_memory
  
      return rc_output, next_m, next_k, next_v
  
imports: import math
  import re
  from functools import partial
  from typing import List, Optional, Tuple
  import torch
  import torch.nn
  from torch import Tensor
  from torch import device
  from fairseq.models import FairseqEncoder
  from fairseq.models.speech_to_text.utils import NoOp, attention_suppression, layer_norm_backward_hook, lengths_to_padding_mask, segments_to_sequence
label_desc: Use image node_id 5 for calling the NoSegAugmentedMemoryMultiheadAttentionBmm obj's underlying member method code with example usage: obj.forward(input, lengths, mems, attention_mask, pre_mems, left_context_key, left_context_val, rpe) and returns: rc_output,  next_m,  next_k,  next_v