class_docstr: Multi-head scaled dot-product attention layer.
  
  Args:
      hidden_size: size of the embeddings in the model, also known as d_model
      num_attention_heads: number of heads in multi-head attention
      attn_score_dropout: probability of dropout applied to attention scores
      attn_layer_dropout: probability of dropout applied to the output of the
          whole layer, but before layer normalization
class_docstr_tok: ["Multi-head","scaled","dot-product","attention","layer",".","Args",":","hidden_size",":","size","of","the","embeddings","in","the","model",",","also","known","as","d_model","num_attention_heads",":","number","of","heads","in","multi-head","attention","attn_score_dropout",":","probability","of","dropout","applied","to","attention","scores","attn_layer_dropout",":","probability","of","dropout","applied","to","the","output","of","the","whole","layer",",","but","before","layer","normalization"]
code: def __init__(
      self,
      hidden_size,
      num_attention_heads,
      attn_score_dropout=0.0,
      attn_layer_dropout=0.0,
  ):
      super().__init__()
      if hidden_size % num_attention_heads != 0:
          raise ValueError(
              "The hidden size (%d) is not a multiple of the number "
              "of attention heads (%d)"
              % (hidden_size, num_attention_heads)
          )
      self.hidden_size = hidden_size
      self.num_attention_heads = num_attention_heads
      self.attn_head_size = int(hidden_size / num_attention_heads)
      self.attn_scale = math.sqrt(math.sqrt(self.attn_head_size))
  
      self.query_net = nn.Linear(hidden_size, hidden_size)
      self.key_net = nn.Linear(hidden_size, hidden_size)
      self.value_net = nn.Linear(hidden_size, hidden_size)
      self.out_projection = nn.Linear(hidden_size, hidden_size)
  
      self.attn_dropout = nn.Dropout(attn_score_dropout)
      self.layer_dropout = nn.Dropout(attn_layer_dropout)
  
imports: import math
  import numpy
  import torch
  from torch import nn
  from torch.nn.functional import gelu
  from nemo.collections.common.parts import form_attention_mask
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(hidden_size, num_attention_heads, attn_score_dropout, attn_layer_dropout)