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,
      query_head_units: Optional[int] = None,
      layout: str = "NKT",
      use_einsum: bool = False,
  ):
      """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."::
  
          .. code-block:: none
  
              score = <h_q, h_k> / sqrt(dim_q)
  
      normalized
          If turned on, the cosine distance is used, i.e::
  
          .. code-block:: none
  
              score = <h_q / ||h_q||, h_k / ||h_k||>
  
      eps
          The epsilon value used in L2 normalization
      query_head_units
          The units of each query head. If it's empty, we will estimate it via the
          shape_array of the query.
      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)
      """
      # TODO(sxjscience) Profile layout
      if normalized:
          query = l2_normalize(query, axis=-1, eps=eps)
          key = l2_normalize(key, axis=-1, eps=eps)
      if scaled:
          if query_head_units is None:
              raise NotImplementedError(
                  "You will need to specify query_head_units!"
              )
          else:
              scale = math.sqrt(query_head_units)
      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 = np.expand_dims(mask, axis=1).astype(np.bool)
          # 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 = npx.batch_dot(query, key, transpose_b=True)
          if edge_scores is not None:
              scores = scores + edge_scores
          attn_weights = masked_softmax(
              scores, mask, axis=-1, temperature=scale
          )
          attn_weights = npx.dropout(attn_weights, p=dropout)
          # 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 = np.einsum(
                  "bnij,bnjc->binc", attn_weights, value
              )
          else:
              context_vec = npx.batch_dot(
                  attn_weights, value
              ).transpose((0, 2, 1, 3))
          context_vec = npx.reshape(context_vec, (-2, -2, -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 = np.expand_dims(mask, axis=1).astype(np.bool)
          # 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 = np.einsum("binc,bjnc->bnij", query, key)
          else:
              scores = npx.batch_dot(
                  np.swapaxes(query, 1, 2),
                  np.swapaxes(key, 1, 2),
                  transpose_b=True,
              )
          if edge_scores is not None:
              scores = scores + edge_scores
          attn_weights = masked_softmax(
              scores, mask, axis=-1, temperature=scale
          )
          attn_weights = npx.dropout(attn_weights, p=dropout)
          # 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 = np.einsum(
                  "bnij,bjnc->binc", attn_weights, value
              )
          else:
              context_vec = npx.batch_dot(
                  attn_weights, np.swapaxes(value, 1, 2)
              ).transpose((0, 2, 1, 3))
          context_vec = npx.reshape(context_vec, (-2, -2, -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 = np.expand_dims(mask, axis=1).astype(np.bool)
          # 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 = np.einsum("ibnc,jbnc->bnij", query, key)
          else:
              scores = npx.batch_dot(
                  query.transpose((1, 2, 0, 3)),
                  key.transpose((1, 2, 3, 0)),
              )
          if edge_scores is not None:
              scores = scores + edge_scores
          attn_weights = masked_softmax(
              scores, mask, axis=-1, temperature=scale
          )
          attn_weights = npx.dropout(attn_weights, p=dropout)
          # 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 = np.einsum(
                  "bnij,jbnc->ibnc", attn_weights, value
              )
          else:
              context_vec = npx.batch_dot(
                  attn_weights, value.transpose((1, 2, 0, 3))
              ).transpose((2, 0, 1, 3))
          context_vec = npx.reshape(context_vec, (-2, -2, -1))
      else:
          raise NotImplementedError(
              'layout="{}" is not supported! '
              'We only support layout = "NKT", "NTK", and "TNK".'.format(
                  layout
              )
          )
      return context_vec, [scores, attn_weights]
  
imports: import math
  import mxnet
  from mxnet import np, npx
  import numpy
  from mxnet.gluon.block import HybridBlock
  from mxnet.gluon import nn
  from .op import l2_normalize
  from .layers import SinusoidalPositionalEmbedding, BucketPositionalEmbedding, LearnedPositionalEmbedding
  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, query_head_units, layout, use_einsum) and returns: context_vec,  list