docstr: :param x: [B, L, Hs]
  :param attention_mask: [B, L] padding后的句子后面补0了，补0的位置为True，前面部分为False
  :param head_mask: [L] [N,L]
  :return:
docstr_tok: [":","param","x",":","[","B",",","L",",","Hs","]",":","param","attention_mask",":","[","B",",","L","]","padding\u540e\u7684\u53e5\u5b50\u540e\u9762\u88650\u4e86\uff0c\u88650\u7684\u4f4d\u7f6e\u4e3aTrue\uff0c\u524d\u9762\u90e8\u5206\u4e3aFalse",":","param","head_mask",":","[","L","]","[","N",",","L","]",":","return",":"]
code: def forward(
      self,
      x,
      key_padding_mask=None,
      attention_mask=None,
      head_mask=None,
  ):
      """
      :param x: [B, L, Hs]
      :param attention_mask: [B, L] padding后的句子后面补0了，补0的位置为True，前面部分为False
      :param head_mask: [L] [N,L]
      :return:
      """
      attention_outputs = self.multihead_attention(
          x, x, x, key_padding_mask, attention_mask, head_mask
      )
      attention_output = attention_outputs[0]
      attention_output = self.dense(attention_output)
      attention_output = self.dropout(attention_output)
      attention_output = self.layerNorm(attention_output + x)
      outputs = (attention_output,) + attention_outputs[
          1:
      ]  # 后面是 attention weight
      return outputs
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
  from .Attention import MultiHeadAttention
label_desc: Use image node_id 2 for calling the TransformerAttention obj's underlying member method code with example usage: obj.forward(x, key_padding_mask, attention_mask, head_mask) and returns: outputs