class_docstr: Multi-Head Attention optimized for efficient ANE deployment
      
docstr: Args:
      embed_dim:          Dimensionality of the input embeddings
      d_qk:               Dimensionality of the query and key embeddings. They must match in order to compute
                          dot product attention. If None, it is set to that of the input tensors, i.e. `embed_dim`
      d_v:                Dimensionality of the value embeddings. It may differ from that of the query and
                          key embeddings. If None, it is set to that of the input tensors.
      d_out:              Dimensionality of the output projection. If None, it is set to that of the input tensors
      n_head:             The number of different attention heads to compute in parallel, uses :math:`d_qk/n_head`
                          channel groups in parallel to learn different attention patterns in a single layer
      dropout:            The probability that each attention weight is zero-masked independent of other weights
docstr_tok: ["Args",":","embed_dim",":","Dimensionality","of","the","input","embeddings","d_qk",":","Dimensionality","of","the","query","and","key","embeddings",".","They","must","match","in","order","to","compute","dot","product","attention",".","If","None",",","it","is","set","to","that","of","the","input","tensors",",","i.e",".","`","embed_dim","`","d_v",":","Dimensionality","of","the","value","embeddings",".","It","may","differ","from","that","of","the","query","and","key","embeddings",".","If","None",",","it","is","set","to","that","of","the","input","tensors",".","d_out",":","Dimensionality","of","the","output","projection",".","If","None",",","it","is","set","to","that","of","the","input","tensors","n_head",":","The","number","of","different","attention","heads","to","compute","in","parallel",",","uses",":","math",":","`","d_qk\/n_head","`","channel","groups","in","parallel","to","learn","different","attention","patterns","in","a","single","layer","dropout",":","The","probability","that","each","attention","weight","is","zero-masked","independent","of","other","weights"]
code: def __init__(
      self,
      embed_dim,
      d_qk=None,
      d_v=None,
      d_out=None,
      n_head=8,
      dropout=0.1,
      **kwargs,
  ):
      """
      Args:
          embed_dim:          Dimensionality of the input embeddings
          d_qk:               Dimensionality of the query and key embeddings. They must match in order to compute
                              dot product attention. If None, it is set to that of the input tensors, i.e. `embed_dim`
          d_v:                Dimensionality of the value embeddings. It may differ from that of the query and
                              key embeddings. If None, it is set to that of the input tensors.
          d_out:              Dimensionality of the output projection. If None, it is set to that of the input tensors
          n_head:             The number of different attention heads to compute in parallel, uses :math:`d_qk/n_head`
                              channel groups in parallel to learn different attention patterns in a single layer
          dropout:            The probability that each attention weight is zero-masked independent of other weights
      """
      super().__init__()
  
      self.d_qk = d_qk or embed_dim
      self.d_v = d_v or embed_dim
      self.d_out = d_out or embed_dim
  
      self.n_head = n_head
      if self.d_qk % self.n_head != 0 or self.d_v % self.n_head != 0:
          raise ValueError(
              f"Either query-key dimensions ({self.d_qk}) or the value embeddings "
              f"dimensions ({self.d_v}) is not divisible by n_head ({self.n_head})"
          )
      self.q_normalize_fact = float(self.d_qk // self.n_head) ** -0.5
  
      self.q_proj = nn.Conv2d(embed_dim, self.d_qk, 1)
      self.v_proj = nn.Conv2d(embed_dim, self.d_v, 1)
      self.k_proj = nn.Conv2d(embed_dim, self.d_qk, 1)
      self.out_proj = nn.Conv2d(self.d_v, self.d_out, 1)
      self.dropout = (
          nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
      )
  
      self.apply(self._reset_parameters)
  
imports: import torch
  import torch.nn
  from .layer_norm import LayerNormANE
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(embed_dim, d_qk, d_v, d_out, n_head, dropout)