class_docstr: The multi-head attention
  
  out = softmax(<Q_i, K_j> + R_{i, j}) V
  
  We support multiple layouts
  
  Let's denote batch_size as B, num_heads as K,
   query_length as L_q, mem_length as L_m, key_dim as C_k, value_dim as C_v
  
  - layout="NKT"
      query: (B, K, L_q, C_k)
      key: (B, K, L_m, C_k)
      value: (B, K, L_m, C_v)
      out: (B, L_q, K * C_v)
  - layout="NTK"
      query: (B, L_q, K, C_k)
      key: (B, L_m, K, C_k)
      value: (B, L_m, K, C_v)
      out: (B, L_q, K * C_v)
  - layout="TNK"
      query: (L_q, B, K, C_k)
      key: (L_m, B, K, C_k)
      value: (L_m, B, K, C_v)
      out: (L_q, B, K * C_v)
class_docstr_tok: ["The","multi-head","attention","out","=","softmax","(","<","Q_i",",","K_j",">","+","R_","{","i",",","j","}",")","V","We","support","multiple","layouts","Let","'s","denote","batch_size","as","B",",","num_heads","as","K",",","query_length","as","L_q",",","mem_length","as","L_m",",","key_dim","as","C_k",",","value_dim","as","C_v","-","layout=","''","NKT","''","query",":","(","B",",","K",",","L_q",",","C_k",")","key",":","(","B",",","K",",","L_m",",","C_k",")","value",":","(","B",",","K",",","L_m",",","C_v",")","out",":","(","B",",","L_q",",","K","*","C_v",")","-","layout=","''","NTK","''","query",":","(","B",",","L_q",",","K",",","C_k",")","key",":","(","B",",","L_m",",","K",",","C_k",")","value",":","(","B",",","L_m",",","K",",","C_v",")","out",":","(","B",",","L_q",",","K","*","C_v",")","-","layout=","''","TNK","''","query",":","(","L_q",",","B",",","K",",","C_k",")","key",":","(","L_m",",","B",",","K",",","C_k",")","value",":","(","L_m",",","B",",","K",",","C_v",")","out",":","(","L_q",",","B",",","K","*","C_v",")"]
code: def __repr__(self):
      s = (
          "{name}(\n"
          "   query_units={query_units},\n"
          "   num_heads={num_heads},\n"
          "   attention_dropout={attention_dropout},\n"
          "   scaled={scaled},\n"
          "   normalized={normalized},\n"
          '   layout="{layout}",\n'
          "   use_einsum={use_einsum}\n"
          ")"
      )
      return s.format(
          name=self.__class__.__name__,
          query_units=self._query_units,
          num_heads=self._num_heads,
          attention_dropout=self._attention_dropout,
          scaled=self._scaled,
          normalized=self._normalized,
          layout=self._layout,
          use_einsum=self._use_einsum,
      )
  
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 4 for calling the MultiHeadAttentionCell obj's underlying member method code with example usage: obj.__repr__() and returns: s