docstr: :param Q: [B, L, Hs]
  :param K: [B, S, Hs]
  :param V: [B, S, Hs]
  :param key_padding_mask: [B, S]                为 1/True 的地方需要 mask
  :param attention_mask: [S] / [L, S] 指定位置 mask 掉， 为 1/True 的地方需要 mask
  :param head_mask: [N] 指定 head mask 掉，        为 1/True 的地方需要 mask
code: def forward(
      self,
      Q,
      K,
      V,
      key_padding_mask=None,
      attention_mask=None,
      head_mask=None,
  ):
      """
      :param Q: [B, L, Hs]
      :param K: [B, S, Hs]
      :param V: [B, S, Hs]
      :param key_padding_mask: [B, S]                为 1/True 的地方需要 mask
      :param attention_mask: [S] / [L, S] 指定位置 mask 掉， 为 1/True 的地方需要 mask
      :param head_mask: [N] 指定 head mask 掉，        为 1/True 的地方需要 mask
      """
      B, L, Hs = Q.shape
      S = V.size(1)
      N, H = self.num_heads, self.head_dim
  
      q = self.q_in(Q).view(B, L, N, H).transpose(1, 2)  # [B, N, L, H]
      k = self.k_in(K).view(B, S, N, H).transpose(1, 2)  # [B, N, S, H]
      v = self.v_in(V).view(B, S, N, H).transpose(1, 2)  # [B, N, S, H]
  
      if key_padding_mask is not None:
          key_padding_mask = key_padding_mask.ne(0)
          key_padding_mask = key_padding_mask.unsqueeze(1).unsqueeze(1)
  
      if attention_mask is not None:
          attention_mask = attention_mask.ne(0)
          if attention_mask.dim() == 1:
              attention_mask = attention_mask.unsqueeze(0)
          elif attention_mask.dim() == 2:
              attention_mask = (
                  attention_mask.unsqueeze(0)
                  .unsqueeze(0)
                  .expand(B, -1, -1, -1)
              )
          else:
              raise ValueError(
                  f"attention_mask dim must be 1 or 2, can not be {attention_mask.dim()}"
              )
  
      if key_padding_mask is None:
          mask_out = (
              attention_mask if attention_mask is not None else None
          )
      else:
          mask_out = (
              (key_padding_mask + attention_mask).ne(0)
              if attention_mask is not None
              else key_padding_mask
          )
  
      if head_mask is not None:
          head_mask = head_mask.eq(0)
          head_mask = head_mask.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
  
      attention_out, attention_weight = self.attention(
          q, k, v, mask_out=mask_out, head_mask=head_mask
      )
  
      attention_out = attention_out.transpose(1, 2).reshape(
          B, L, N * H
      )  # [B, N, L, H] -> [B, L, N * H]
  
      # concat all heads, and do output linear
      attention_out = self.out(
          attention_out
      )  # [B, L, N * H] -> [B, L, H]
  
      if self.output_attentions:
          return attention_out, attention_weight
      else:
          return (attention_out,)
  
file_path: DeepKE/src/deepke/relation_extraction/standard/module/Attention.py
imports: import logging
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(Q, K, V, key_padding_mask, attention_mask, head_mask) and returns: attention_out,  attention_weight, attention_out