class_docstr: Base class for Multi-head attention.
  
  Parameters
  ----------
  att_dim_in
      Attention dimension (number of hidden units)
  heads
      Number of attention heads
  att_dim_out
      Output dimension (number of output units)
  dropout
      Dropout rate on attention scores
class_docstr_tok: ["Base","class","for","Multi-head","attention",".","Parameters","--","--","--","--","--","att_dim_in","Attention","dimension","(","number","of","hidden","units",")","heads","Number","of","attention","heads","att_dim_out","Output","dimension","(","number","of","output","units",")","dropout","Dropout","rate","on","attention","scores"]
docstr: Returns context vectors of multi-head dot attention.
  
  Parameters
  ----------
  queries
      Queries tensor of shape (batch_size, query_max_length, dim)
  keys
      Keys tensor of shape (batch_size, memory_max_length, dim)
  values
      Values tensor of shape (batch_size, memory_max_length, dim)
  mask
  
  Returns
  -------
  Context vectors of shape (batch_size, query_max_length, att_dim_out)
docstr_tok: ["Returns","context","vectors","of","multi-head","dot","attention",".","Parameters","--","--","--","--","--","queries","Queries","tensor","of","shape","(","batch_size",",","query_max_length",",","dim",")","keys","Keys","tensor","of","shape","(","batch_size",",","memory_max_length",",","dim",")","values","Values","tensor","of","shape","(","batch_size",",","memory_max_length",",","dim",")","mask","Returns","--","--","--","-","Context","vectors","of","shape","(","batch_size",",","query_max_length",",","att_dim_out",")"]
code: def _attend(
      self,
      F,
      queries: Tensor,
      keys: Tensor,
      values: Tensor,
      mask: Optional[Tensor] = None,
  ) -> Tensor:
      r"""
      Returns context vectors of multi-head dot attention.
  
      Parameters
      ----------
      queries
          Queries tensor of shape (batch_size, query_max_length, dim)
      keys
          Keys tensor of shape (batch_size, memory_max_length, dim)
      values
          Values tensor of shape (batch_size, memory_max_length, dim)
      mask
  
      Returns
      -------
      Context vectors of shape (batch_size, query_max_length, att_dim_out)
      """
  
      # scale by 1/sqrt(dim_per_head)
      queries = queries * (self.dim_per_head**-0.5)
  
      # (batch_size * heads, length, dim/heads)
      queries = split_heads(F, queries, self.dim_per_head, self.heads)
      keys = split_heads(F, keys, self.dim_per_head, self.heads)
      values = split_heads(F, values, self.dim_per_head, self.heads)
  
      # (batch_size * heads, query_max_length, dim_per_head)
      contexts = dot_attention(
          F, queries, keys, values, mask=mask, dropout=self.dropout
      )
  
      # (batch_size, query_max_length, input_dim)
      contexts = combine_heads(
          F, contexts, self.dim_per_head, self.heads
      )
  
      # contexts: (batch_size, query_max_length, output_dim)
      contexts = self.dense_att(contexts)
  
      return contexts
  
imports: from typing import Dict, Optional, Tuple
  import mxnet
  from mxnet.gluon import HybridBlock
  from gluonts.mx import Tensor
label_desc: Use image node_id 2 for calling the MultiHeadAttentionBase obj's underlying member method code with example usage: obj._attend(F, queries, keys, values, mask) and returns: contexts