docstr: Args:
      query, key, value (Tensor): map a query and a set of key-value pairs to an output.
          See "Attention Is All You Need" for more details.
      attn_mask, bias_k and bias_v (Tensor, optional): keyword arguments passed to the attention layer.
          See the definitions in the attention.
  Shape:
      - Inputs:
      - query: :math:`(L, N, E)`
      - key: :math:`(S, N, E)`
      - value: :math:`(S, N, E)`
      - attn_mask, bias_k and bias_v: same with the shape of the corresponding args in attention layer.
      - Outputs:
      - attn_output: :math:`(L, N, E)`
      - attn_output_weights: :math:`(N * H, L, S)`
      where where L is the target length, S is the sequence length, H is the number of attention heads,
          N is the batch size, and E is the embedding dimension.
docstr_tok: ["Args",":","query",",","key",",","value","(","Tensor",")",":","map","a","query","and","a","set","of","key-value","pairs","to","an","output",".","See","``","Attention","Is","All","You","Need","''","for","more","details",".","attn_mask",",","bias_k","and","bias_v","(","Tensor",",","optional",")",":","keyword","arguments","passed","to","the","attention","layer",".","See","the","definitions","in","the","attention",".","Shape",":","-","Inputs",":","-","query",":",":","math",":","`","(","L",",","N",",","E",")","`","-","key",":",":","math",":","`","(","S",",","N",",","E",")","`","-","value",":",":","math",":","`","(","S",",","N",",","E",")","`","-","attn_mask",",","bias_k","and","bias_v",":","same","with","the","shape","of","the","corresponding","args","in","attention","layer",".","-","Outputs",":","-","attn_output",":",":","math",":","`","(","L",",","N",",","E",")","`","-","attn_output_weights",":",":","math",":","`","(","N","*","H",",","L",",","S",")","`","where","where","L","is","the","target","length",",","S","is","the","sequence","length",",","H","is","the","number","of","attention","heads",",","N","is","the","batch","size",",","and","E","is","the","embedding","dimension","."]
code: def forward(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      attn_mask: Optional[torch.Tensor] = None,
      bias_k: Optional[torch.Tensor] = None,
      bias_v: Optional[torch.Tensor] = None,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      r"""
      Args:
          query, key, value (Tensor): map a query and a set of key-value pairs to an output.
              See "Attention Is All You Need" for more details.
          attn_mask, bias_k and bias_v (Tensor, optional): keyword arguments passed to the attention layer.
              See the definitions in the attention.
      Shape:
          - Inputs:
          - query: :math:`(L, N, E)`
          - key: :math:`(S, N, E)`
          - value: :math:`(S, N, E)`
          - attn_mask, bias_k and bias_v: same with the shape of the corresponding args in attention layer.
          - Outputs:
          - attn_output: :math:`(L, N, E)`
          - attn_output_weights: :math:`(N * H, L, S)`
          where where L is the target length, S is the sequence length, H is the number of attention heads,
              N is the batch size, and E is the embedding dimension.
      """
      tgt_len, src_len, bsz, embed_dim = (
          query.size(-3),
          key.size(-3),
          query.size(-2),
          query.size(-1),
      )
      q, k, v = self.in_proj_container(query, key, value)
      assert (
          q.size(-1) % self.nhead == 0
      ), "query's embed_dim must be divisible by the number of heads"
      head_dim = q.size(-1) // self.nhead
      q = q.reshape(tgt_len, bsz * self.nhead, head_dim)
  
      assert (
          k.size(-1) % self.nhead == 0
      ), "key's embed_dim must be divisible by the number of heads"
      head_dim = k.size(-1) // self.nhead
      k = k.reshape(src_len, bsz * self.nhead, head_dim)
  
      assert (
          v.size(-1) % self.nhead == 0
      ), "value's embed_dim must be divisible by the number of heads"
      head_dim = v.size(-1) // self.nhead
      v = v.reshape(src_len, bsz * self.nhead, head_dim)
  
      attn_output, attn_output_weights = self.attention_layer(
          q, k, v, attn_mask=attn_mask, bias_k=bias_k, bias_v=bias_v
      )
      attn_output = attn_output.reshape(tgt_len, bsz, embed_dim)
      attn_output = self.out_proj(attn_output)
      return attn_output, attn_output_weights
  
imports: import math
  from collections import OrderedDict
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
label_desc: Use image node_id 2 for calling the MultiheadAttentionContainer obj's underlying member method code with example usage: obj.forward(query, key, value, attn_mask, bias_k, bias_v) and returns: attn_output,  attn_output_weights