docstr: A multi-head attention container
  Args:
      nhead: the number of heads in the multiheadattention model
      in_proj_container: A container of multi-head in-projection linear layers (a.k.a nn.Linear).
      attention_layer: The attention layer.
      out_proj: The multi-head out-projection layer (a.k.a nn.Linear).
  Examples::
      >>> import torch
      >>> embed_dim, num_heads, bsz = 10, 5, 64
      >>> in_proj_container = InProjContainer(torch.nn.Linear(embed_dim, embed_dim),
                                              torch.nn.Linear(embed_dim, embed_dim),
                                              torch.nn.Linear(embed_dim, embed_dim))
      >>> MHA = MultiheadAttentionContainer(num_heads,
                                            in_proj_container,
                                            ScaledDotProduct(),
                                            torch.nn.Linear(embed_dim, embed_dim))
      >>> query = torch.rand((21, bsz, embed_dim))
      >>> key = value = torch.rand((16, bsz, embed_dim))
      >>> attn_output, attn_weights = MHA(query, key, value)
      >>> print(attn_output.shape)
      >>> torch.Size([21, 64, 10])
docstr_tok: ["A","multi-head","attention","container","Args",":","nhead",":","the","number","of","heads","in","the","multiheadattention","model","in_proj_container",":","A","container","of","multi-head","in-projection","linear","layers","(","a.k.a","nn.Linear",")",".","attention_layer",":","The","attention","layer",".","out_proj",":","The","multi-head","out-projection","layer","(","a.k.a","nn.Linear",")",".","Examples",":",":",">",">",">","import","torch",">",">",">","embed_dim",",","num_heads",",","bsz","=","10",",","5",",","64",">",">",">","in_proj_container","=","InProjContainer","(","torch.nn.Linear","(","embed_dim",",","embed_dim",")",",","torch.nn.Linear","(","embed_dim",",","embed_dim",")",",","torch.nn.Linear","(","embed_dim",",","embed_dim",")",")",">",">",">","MHA","=","MultiheadAttentionContainer","(","num_heads",",","in_proj_container",",","ScaledDotProduct","(",")",",","torch.nn.Linear","(","embed_dim",",","embed_dim",")",")",">",">",">","query","=","torch.rand","(","(","21",",","bsz",",","embed_dim",")",")",">",">",">","key","=","value","=","torch.rand","(","(","16",",","bsz",",","embed_dim",")",")",">",">",">","attn_output",",","attn_weights","=","MHA","(","query",",","key",",","value",")",">",">",">","print","(","attn_output.shape",")",">",">",">","torch.Size","(","[","21",",","64",",","10","]",")"]
code: def __init__(
      self, nhead, in_proj_container, attention_layer, out_proj
  ):
      r"""A multi-head attention container
      Args:
          nhead: the number of heads in the multiheadattention model
          in_proj_container: A container of multi-head in-projection linear layers (a.k.a nn.Linear).
          attention_layer: The attention layer.
          out_proj: The multi-head out-projection layer (a.k.a nn.Linear).
      Examples::
          >>> import torch
          >>> embed_dim, num_heads, bsz = 10, 5, 64
          >>> in_proj_container = InProjContainer(torch.nn.Linear(embed_dim, embed_dim),
                                                  torch.nn.Linear(embed_dim, embed_dim),
                                                  torch.nn.Linear(embed_dim, embed_dim))
          >>> MHA = MultiheadAttentionContainer(num_heads,
                                                in_proj_container,
                                                ScaledDotProduct(),
                                                torch.nn.Linear(embed_dim, embed_dim))
          >>> query = torch.rand((21, bsz, embed_dim))
          >>> key = value = torch.rand((16, bsz, embed_dim))
          >>> attn_output, attn_weights = MHA(query, key, value)
          >>> print(attn_output.shape)
          >>> torch.Size([21, 64, 10])
      """
      super().__init__()
      self.nhead = nhead
      self.in_proj_container = in_proj_container
      self.attention_layer = attention_layer
      self.out_proj = out_proj
  
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 1 to create a new MultiheadAttentionContainer object from inherited base classes: torch.nn with example: obj = MultiheadAttentionContainer(nhead, in_proj_container, attention_layer, out_proj)