code: def multi_head_dot_attn(
      query,
      key,
      value,
      mask=None,
      edge_scores=None,
      dropout: float = 0.0,
      scaled: bool = True,
      normalized: bool = False,
      eps: float = 1e-6,
      layout: str = "NKT",
      use_einsum: bool = None,
      *,
      training: bool = True,
  ):
      """Multihead dot product attention between the query, key, value.
  
      scaled is False, normalized is False:
          D(h_q, h_k) = <h_q, h_k>
      scaled is True, normalized is False:
          D(h_q, h_k) = <h_q, h_k> / sqrt(dim_q)
      scaled is False, normalized is True:
          D(h_q, h_k) = <h_q / ||h_q||, h_k / ||h_k||>
      scaled is True, normalized is True:
          D(h_q, h_k) = <h_q / ||h_q||, h_k / ||h_k||> / sqrt(dim_q)
  
      If edge_scores is provided, we will calcualte the attention as
          scores = D(h_q, h_k) + EdgeScore_{q, k}
  
      Parameters
      ----------
      query
          Query. The shape depends on the layout
          - layout is 'NKT'
              Shape (batch_size, num_heads, query_length, key_dim)
          - layout is 'NTK'
              Shape (batch_size, query_length, num_heads, key_dim)
          - layout is 'TNK'
              Shape (query_length, batch_size, num_heads, key_dim)
      key
          Key. The shape depends on the layout
          - layout is 'NKT'
              Shape (batch_size, num_heads, mem_length, key_dim)
          - layout is 'NTK'
              Shape (batch_size, mem_length, num_heads, key_dim)
          - layout is 'TNK'
              Shape (mem_length, batch_size, num_heads, key_dim)
      value
          Value. The shape depends on the layout
          - layout is 'NKT'
              Shape (batch_size, num_heads, mem_length, value_dim)
          - layout is 'NTK'
              Shape (batch_size, mem_length, num_heads, value_dim)
          - layout is 'TNK'
              Shape (mem_length, batch_size, num_heads, value_dim)
      mask
          Mask between query and memory. Shape (batch_size, query_length, mem_length)
      edge_scores
          The edge attention score. Shape can be any shape that is broadcastable to
          (batch_size, num_heads, query_length, mem_length)
      dropout
          Dropout rate
      scaled
          Whether to divide the attention weights by the sqrt of the query dimension.
          This is first proposed in "[NIPS2017] Attention is all you need."::
  
              score = <h_q, h_k> / sqrt(dim_q)
  
      normalized
          If turned on, the cosine distance is used, i.e::
  
              score = <h_q / ||h_q||, h_k / ||h_k||>
  
      eps
          The epsilon value used in L2 normalization
      layout
          This stands for the layout of the attention cell. The shape of the input/output will depend
          on the layout. Currently, we support 'NKT', 'NTK' and 'TNK' in which
          'N' means the batch_size, 'K' means the head, and 'T' means the length dimension.
      use_einsum
          Whether to use einsum for the computation
  
      Returns
      -------
      context_vec
          - layout is 'NKT' or 'NTK'
              Shape (batch_size, query_length, num_heads * value_units)
          - layout is 'TNK'
              Shape (query_length, batch_size, num_heads * value_units)
      additional_info
          scores:
              Shape (batch_size, num_head, query_length, mem_length)
          attn_weight:
              Shape (batch_size, num_head, query_length, mem_length)
      """
      if use_einsum is None:
          use_einsum = use_einsum_optimization()
      # TODO(sxjscience) Profile layout
      if normalized:
          query = F.normalize(query, p=2, dim=-1, eps=eps)
          key = F.normalize(key, p=2, dim=-1, eps=eps)
  
      if scaled:
          scale = math.sqrt(query.shape[-1])
      else:
          scale = None
  
      if layout == "NKT":
          # 1. Expand the dimension of the mask:
          #   (B, L_query, L_mem) --> (B, 1, L_query, L_mem)
          if mask is not None:
              mask = th.unsqueeze(mask, dim=1)
          # 2. Calculate the attention weights
          #   Score: (B, N, L_query, C_Q) X (B, N, L_mem, C_Q) --> (B, N, L_query, L_mem)
          scores = th.matmul(query, th.transpose(key, -2, -1))
  
          if edge_scores is not None:
              scores = scores + edge_scores
          attn_weights = masked_softmax(
              scores / scale if scale is not None else scores,
              mask,
              axis=-1,
          )
          attn_weights = th.nn.functional.dropout(
              attn_weights, p=dropout, training=training
          )
          # 3. Calculate the context vector
          # (B, N, L_query, L_mem) X (B, N, L_mem, C_V) --> (B, L_query, N * C_V)
          if use_einsum:
              context_vec = th.einsum(
                  "bnij,bnjc->binc", attn_weights, value
              )
          else:
              context_vec = th.transpose(
                  th.matmul(attn_weights, value), 1, 2
              )
          context_vec = th.reshape(
              context_vec,
              (context_vec.shape[0], context_vec.shape[1], -1),
          )
      elif layout == "NTK":
          # 1. Expand the dimension of the mask:
          #   (B, L_query, L_mem) --> (B, 1, L_query, L_mem)
          if mask is not None:
              mask = th.unsqueeze(mask, dim=1)
          # 2. Calculate the attention weights
          #   Score: (B, L_query, N, C_Q) X (B, L_mem, N, C_Q) --> (B, N, L_query, L_mem)
          if use_einsum:
              scores = th.einsum("binc,bjnc->bnij", query, key)
          else:
              scores = th.matmul(
                  th.transpose(query, 1, 2), key.permute(0, 2, 3, 1)
              )
          if edge_scores is not None:
              scores = scores + edge_scores
          attn_weights = masked_softmax(
              scores / scale if scale is not None else scores, mask
          )
          attn_weights = th.nn.functional.dropout(
              attn_weights, p=dropout, training=training
          )
          # 3. Calculate the context vector
          # (B, N, L_query, L_mem) X (B, L_mem, N, C_V) --> (B, L_query, N * C_V)
          if use_einsum:
              context_vec = th.einsum(
                  "bnij,bjnc->binc", attn_weights, value
              )
          else:
              context_vec = th.matmul(
                  attn_weights, th.transpose(value, 1, 2)
              ).permute(0, 2, 1, 3)
          context_vec = th.reshape(
              context_vec,
              (context_vec.shape[0], context_vec.shape[1], -1),
          )
      elif layout == "TNK":
          # 1. Expand the dimension of the mask:
          #   (B, L_query, L_mem) --> (B, 1, L_query, L_mem)
          if mask is not None:
              mask = th.unsqueeze(mask, dim=1)
          # 2. Calculate the attention weights
          #   Score: (L_query, B, N, C_Q) X (L_mem, B, N, C_Q) --> (B, N, L_query, L_mem)
          #   This layout structure can be implemented very efficiently because B, N are consecutive
          #   to each other. To have a clear picture of what's happening, we may consider the
          #   (i, j)th element of the output
          #       out[i, j, :, :] = query[:, i, j, :] X key[:, i, j, :].T, which is just one GEMM call
          #   We can thus implement the whole kernel via a single call of batched GEMM with stride.
          if use_einsum:
              scores = th.einsum("ibnc,jbnc->bnij", query, key)
          else:
              scores = th.matmul(
                  query.permute(1, 2, 0, 3), key.permute(1, 2, 3, 0)
              )
          if edge_scores is not None:
              scores = scores + edge_scores
          attn_weights = masked_softmax(
              scores / scale if scale is not None else scores, mask
          )
          attn_weights = th.nn.functional.dropout(
              attn_weights, p=dropout, training=training
          )
          # 3. Calculate the context vector
          # (B, N, L_query, L_mem) X (L_mem, B, N, C_V) --> (L_query, B, N * C_V)
          # Again, we can implement it via a single call to batched GEMM with stride.
  
          # Shape (B, N, L_query, C_V)
          if use_einsum:
              context_vec = th.einsum(
                  "bnij,jbnc->ibnc", attn_weights, value
              )
          else:
              context_vec = th.matmul(
                  attn_weights, value.permute(1, 2, 0, 3)
              ).permute(2, 0, 1, 3)
          context_vec = th.reshape(
              context_vec,
              (context_vec.shape[0], context_vec.shape[1], -1),
          )
      else:
          raise NotImplementedError(
              'layout="{}" is not supported! '
              'We only support layout = "NKT", "NTK", and "TNK".'.format(
                  layout
              )
          )
      return context_vec, [scores, attn_weights]
  
file_path: gluon-nlp/src/gluonnlp/torch/attention_cell.py
imports: import math
  import torch
  from torch import nn
  import torch.nn.functional
  from ..base import use_einsum_optimization
  from typing import Optional
label_desc: Use image node_id 5 for calling a global function with example usage: multi_head_dot_attn(query, key, value, mask, edge_scores, dropout, scaled, normalized, eps, layout, use_einsum) and returns: context_vec,  list