parameters: (self, nhead, in_proj_container, attention_layer, out_proj)
docstring: 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])
docstring_summary: 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])
docstring_tokens: 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 ] )
function: 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(MultiheadAttentionContainer, self).__init__()
          self.nhead = nhead
          self.in_proj_container = in_proj_container
          self.attention_layer = attention_layer
          self.out_proj = out_proj
function_tokens: def __init__ ( self , nhead , in_proj_container , attention_layer , out_proj ) : super ( MultiheadAttentionContainer , self ) . __init__ ( ) self . nhead = nhead self . in_proj_container = in_proj_container self . attention_layer = attention_layer self . out_proj = out_proj