class_docstr: Multi-Head Attention layer of Transformer.
  Args:
      n_head (int): number of heads
      n_feat (int): size of the features
      dropout_rate (float): dropout rate
class_docstr_tok: ["Multi-Head","Attention","layer","of","Transformer",".","Args",":","n_head","(","int",")",":","number","of","heads","n_feat","(","int",")",":","size","of","the","features","dropout_rate","(","float",")",":","dropout","rate"]
docstr: Compute 'Scaled Dot Product Attention'.
  Args:
      query (torch.Tensor): (batch, time1, size)
      key (torch.Tensor): (batch, time2, size)
      value(torch.Tensor): (batch, time2, size)
      mask (torch.Tensor): (batch, time1, time2)
      cache (torch.Tensor) : (batch, time_cache, size)
  
  returns:
      output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
      cache (torch.Tensor) : (batch, time_cache_next, size)
docstr_tok: ["Compute","'Scaled","Dot","Product","Attention","'",".","Args",":","query","(","torch.Tensor",")",":","(","batch",",","time1",",","size",")","key","(","torch.Tensor",")",":","(","batch",",","time2",",","size",")","value","(","torch.Tensor",")",":","(","batch",",","time2",",","size",")","mask","(","torch.Tensor",")",":","(","batch",",","time1",",","time2",")","cache","(","torch.Tensor",")",":","(","batch",",","time_cache",",","size",")","returns",":","output","(","torch.Tensor",")",":","transformed","`","value","`","(","batch",",","time1",",","d_model",")","weighted","by","the","query","dot","key","attention","cache","(","torch.Tensor",")",":","(","batch",",","time_cache_next",",","size",")"]
code: def forward(self, query, key, value, mask, pos_emb=None, cache=None):
      """Compute 'Scaled Dot Product Attention'.
      Args:
          query (torch.Tensor): (batch, time1, size)
          key (torch.Tensor): (batch, time2, size)
          value(torch.Tensor): (batch, time2, size)
          mask (torch.Tensor): (batch, time1, time2)
          cache (torch.Tensor) : (batch, time_cache, size)
  
      returns:
          output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
          cache (torch.Tensor) : (batch, time_cache_next, size)
      """
      key, value, query, cache = self.update_cache(
          key=key, value=value, query=query, cache=cache
      )
  
      if torch.is_autocast_enabled():
          query, key, value = (
              query.to(torch.float32),
              key.to(torch.float32),
              value.to(torch.float32),
          )
  
      # temporary until we solve this more gracefully
      with avoid_float16_autocast_context():
          q, k, v = self.forward_qkv(query, key, value)
          scores = torch.matmul(q, k.transpose(-2, -1)) / self.s_d_k
          out = self.forward_attention(v, scores, mask)
      if cache is None:
          return out
      else:
          return out, cache
  
imports: import math
  from functools import lru_cache
  from typing import List, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from nemo.utils import avoid_float16_autocast_context
label_desc: Use image node_id 4 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask, pos_emb, cache) and returns: out, out,  cache