class_docstr: `Attention is all you need <https://arxiv.org/abs/1706.03762>`_ 中提到的多头注意力
  
  :param d_model:
  :param n_head:
  :param dropout:
  :param layer_idx:
class_docstr_tok: ["`","Attention","is","all","you","need","<","https",":","\/\/arxiv.org\/abs\/1706.03762",">","`","_","\u4e2d\u63d0\u5230\u7684\u591a\u5934\u6ce8\u610f\u529b",":","param","d_model",":",":","param","n_head",":",":","param","dropout",":",":","param","layer_idx",":"]
docstr: :param query: ``[batch, seq, dim]``
  :param key: ``[batch, seq, dim]``
  :param value: ``[batch, seq, dim]``
  :param key_mask: ``[batch, seq]`` 用于指示哪些 ``key`` 不要 attend 到；注意到 mask 为 **1** 的地方是要attend到的
  :param attn_mask: ``[seq, seq]``, 用于 mask 掉 attention map。 主要是用在训练时 decoder 端的 :class:`SelfAttention` ，
      下三角为 1。
  :param state: 过去的信息，在 inference 的时候会用到，比如 encoder output、decoder 的 prev kv。这样可以减少计算。
  :return:
docstr_tok: [":","param","query",":","``","[","batch",",","seq",",","dim","]","``",":","param","key",":","``","[","batch",",","seq",",","dim","]","``",":","param","value",":","``","[","batch",",","seq",",","dim","]","``",":","param","key_mask",":","``","[","batch",",","seq","]","``","\u7528\u4e8e\u6307\u793a\u54ea\u4e9b","``","key","``","\u4e0d\u8981","attend","\u5230\uff1b\u6ce8\u610f\u5230","mask","\u4e3a","*","*","1","*","*","\u7684\u5730\u65b9\u662f\u8981attend\u5230\u7684",":","param","attn_mask",":","``","[","seq",",","seq","]","``",",","\u7528\u4e8e","mask","\u6389","attention","map\u3002","\u4e3b\u8981\u662f\u7528\u5728\u8bad\u7ec3\u65f6","decoder","\u7aef\u7684",":","class",":","`","SelfAttention","`","\uff0c","\u4e0b\u4e09\u89d2\u4e3a","1\u3002",":","param","state",":","\u8fc7\u53bb\u7684\u4fe1\u606f\uff0c\u5728","inference","\u7684\u65f6\u5019\u4f1a\u7528\u5230\uff0c\u6bd4\u5982","encoder","output\u3001decoder","\u7684","prev","kv\u3002\u8fd9\u6837\u53ef\u4ee5\u51cf\u5c11\u8ba1\u7b97\u3002",":","return",":"]
code: def forward(
      self, query, key, value, key_mask=None, attn_mask=None, state=None
  ):
      """
  
      :param query: ``[batch, seq, dim]``
      :param key: ``[batch, seq, dim]``
      :param value: ``[batch, seq, dim]``
      :param key_mask: ``[batch, seq]`` 用于指示哪些 ``key`` 不要 attend 到；注意到 mask 为 **1** 的地方是要attend到的
      :param attn_mask: ``[seq, seq]``, 用于 mask 掉 attention map。 主要是用在训练时 decoder 端的 :class:`SelfAttention` ，
          下三角为 1。
      :param state: 过去的信息，在 inference 的时候会用到，比如 encoder output、decoder 的 prev kv。这样可以减少计算。
      :return:
      """
      assert key.size() == value.size()
      if state is not None:
          assert self.layer_idx is not None
      qkv_same = query.data_ptr() == key.data_ptr() == value.data_ptr()
  
      q = self.q_proj(query)  # batch x seq x dim
      q *= self.scaling
      k = v = None
      prev_k = prev_v = None
  
      # 从state中取kv
      if isinstance(state, TransformerState):  # 说明此时在inference阶段
          if qkv_same:  # 此时在decoder self attention
              prev_k = state.decoder_prev_key[self.layer_idx]
              prev_v = state.decoder_prev_value[self.layer_idx]
          else:  # 此时在decoder-encoder attention，直接将保存下来的key装载起来即可
              k = state.encoder_key[self.layer_idx]
              v = state.encoder_value[self.layer_idx]
  
      if k is None:
          k = self.k_proj(key)
          v = self.v_proj(value)
  
      if prev_k is not None:
          k = torch.cat((prev_k, k), dim=1)
          v = torch.cat((prev_v, v), dim=1)
  
      # 更新state
      if isinstance(state, TransformerState):
          if qkv_same:
              state.decoder_prev_key[self.layer_idx] = k
              state.decoder_prev_value[self.layer_idx] = v
          else:
              state.encoder_key[self.layer_idx] = k
              state.encoder_value[self.layer_idx] = v
  
      # 开始计算attention
      batch_size, q_len, d_model = query.size()
      k_len, v_len = k.size(1), v.size(1)
      q = q.reshape(batch_size, q_len, self.n_head, self.head_dim)
      k = k.reshape(batch_size, k_len, self.n_head, self.head_dim)
      v = v.reshape(batch_size, v_len, self.n_head, self.head_dim)
  
      attn_weights = torch.einsum(
          "bqnh,bknh->bqkn", q, k
      )  # bs,q_len,k_len,n_head
      if key_mask is not None:
          _key_mask = ~key_mask[
              :, None, :, None
          ].bool()  # batch,1,k_len,1
          attn_weights = attn_weights.masked_fill(
              _key_mask, -float("inf")
          )
  
      if attn_mask is not None:
          _attn_mask = attn_mask[None, :, :, None].eq(
              0
          )  # 1,q_len,k_len,n_head
          attn_weights = attn_weights.masked_fill(
              _attn_mask, -float("inf")
          )
  
      attn_weights = F.softmax(attn_weights, dim=2)
      attn_weights = F.dropout(
          attn_weights, p=self.dropout, training=self.training
      )
  
      output = torch.einsum(
          "bqkn,bknh->bqnh", attn_weights, v
      )  # batch,q_len,n_head,head_dim
      output = output.reshape(batch_size, q_len, -1)
      output = self.out_proj(output)  # batch,q_len,dim
  
      return output, attn_weights
  
imports: import math
  import torch
  import torch.nn.functional
  from torch import nn
  from .decoder.seq2seq_state import TransformerState
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, key_mask, attn_mask, state) and returns: output,  attn_weights