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 forward(self, queries, keys, values, attention_mask):
      # attention_mask is needed to hide the tokens which correspond to [PAD]
      # in the case of BERT, or to hide the future tokens in the case of
      # vanilla language modeling and translation
      query = self.query_net(queries)
      key = self.key_net(keys)
      value = self.value_net(values)
      query = self.transpose_for_scores(query) / self.attn_scale
      key = self.transpose_for_scores(key) / self.attn_scale
      value = self.transpose_for_scores(value)
  
      # for numerical stability we pre-divide query and key by sqrt(sqrt(d))
      attention_scores = torch.matmul(query, key.transpose(-1, -2))
      if attention_mask is not None:
          attention_scores = attention_scores + attention_mask.to(
              attention_scores.dtype
          )
      attention_probs = torch.softmax(attention_scores, dim=-1)
      attention_probs = self.attn_dropout(attention_probs)
  
      context = torch.matmul(attention_probs, value)
      context = context.permute(0, 2, 1, 3).contiguous()
      new_context_shape = context.size()[:-2] + (self.hidden_size,)
      context = context.view(*new_context_shape)
  
      # output projection
      output_states = self.out_projection(context)
      output_states = self.layer_dropout(output_states)
      return output_states
  
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 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(queries, keys, values, attention_mask) and returns: output_states