class_docstr: Get the score based on the query and relative position index. This is used for implementing
   relative attention.
  
  For the multi-head attention with relative positional encoding, we have the formula
  
  .. math::
  
      out = \text{softmax}(\frac{Q K^T + R}{\sqrt{d}}) V
  
  
  Here, :math:`R` is the relative positional encoding matrix.
  
  Usually, :math:`R_{i, j}` is calculate based on the
  relative positional difference :math:`i - j`.
  
  This function aims at generating the R matrix given the query and the relative positions.
  We support the following methods:
  
  - method = 'transformer_xl'
      :math:`R_{i, j} = <Q, W S_{i - j}>`, in which :math:`S_{i, j}` is the sinusoidal embedding and
      :math:`W` is a Dense layer that maps :math:`S_{i - j}` to the same dimension as the query.
      This is proposed in paper:
  
          [ACL2019] Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context
  
  - method = 'shaw'
      :math:`R_{i, j} = < Q, E_{i - j}>`, in which :math:`E_{i - j}` is the learned positional embedding
      This is proposed in paper:
  
          [NAACL2018] Self-Attention with Relative Position Representations
  
  - method = 't5'
      :math:`R_{i, j} = E_{i - j}`, in which :math:`E_{i - j}` is the bucket positional embedding.
      This is proposed in paper:
  
          [Arxiv2019] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
  
  Like in MultiHeadAttentionCell, we support different layouts to cope with the query matrix.
  
  - layout="NKT"
      query: (B, K, L_q, C_k)
  - layout="NTK"
      query: (B, L_q, K, C_k)
  - layout="TNK"
      query: (L_q, B, K, C_k)
class_docstr_tok: ["Get","the","score","based","on","the","query","and","relative","position","index",".","This","is","used","for","implementing","relative","attention",".","For","the","multi-head","attention","with","relative","positional","encoding",",","we","have","the","formula","..","math",":",":","out","=","\\text","{","softmax","}","(","\\frac","{","Q","K^T","+","R","}","{","\\sqrt","{","d","}","}",")","V","Here",",",":","math",":","`","R","`","is","the","relative","positional","encoding","matrix",".","Usually",",",":","math",":","`","R_","{","i",",","j","}","`","is","calculate","based","on","the","relative","positional","difference",":","math",":","`","i","-","j","`",".","This","function","aims","at","generating","the","R","matrix","given","the","query","and","the","relative","positions",".","We","support","the","following","methods",":","-","method","=","'transformer_xl'",":","math",":","`","R_","{","i",",","j","}","=","<","Q",",","W","S_","{","i","-","j","}",">","`",",","in","which",":","math",":","`","S_","{","i",",","j","}","`","is","the","sinusoidal","embedding","and",":","math",":","`","W","`","is","a","Dense","layer","that","maps",":","math",":","`","S_","{","i","-","j","}","`","to","the","same","dimension","as","the","query",".","This","is","proposed","in","paper",":","[","ACL2019","]","Transformer-XL",":","Attentive","Language","Models","Beyond","a","Fixed-Length","Context","-","method","=","'shaw'",":","math",":","`","R_","{","i",",","j","}","=","<","Q",",","E_","{","i","-","j","}",">","`",",","in","which",":","math",":","`","E_","{","i","-","j","}","`","is","the","learned","positional","embedding","This","is","proposed","in","paper",":","[","NAACL2018","]","Self-Attention","with","Relative","Position","Representations","-","method","=","'t5'",":","math",":","`","R_","{","i",",","j","}","=","E_","{","i","-","j","}","`",",","in","which",":","math",":","`","E_","{","i","-","j","}","`","is","the","bucket","positional","embedding",".","This","is","proposed","in","paper",":","[","Arxiv2019","]","Exploring","the","Limits","of","Transfer","Learning","with","a","Unified","Text-to-Text","Transformer","Like","in","MultiHeadAttentionCell",",","we","support","different","layouts","to","cope","with","the","query","matrix",".","-","layout=","''","NKT","''","query",":","(","B",",","K",",","L_q",",","C_k",")","-","layout=","''","NTK","''","query",":","(","B",",","L_q",",","K",",","C_k",")","-","layout=","''","TNK","''","query",":","(","L_q",",","B",",","K",",","C_k",")"]
docstr: Parameters
  ----------
  query_units
  num_heads
  pos_embed_units
  max_distance
  bidirectional
  num_buckets
  method
  dropout
  dtype
  layout
  use_einsum
docstr_tok: ["Parameters","--","--","--","--","--","query_units","num_heads","pos_embed_units","max_distance","bidirectional","num_buckets","method","dropout","dtype","layout","use_einsum"]
code: def __init__(
      self,
      query_units,
      num_heads,
      pos_embed_units: Optional[int] = None,
      max_distance=None,
      bidirectional=False,
      num_buckets=None,
      method="transformer_xl",
      dropout: float = 0.0,
      dtype="float32",
      layout="NTK",
      use_einsum=False,
      embed_initializer=None,
  ):
      """
  
      Parameters
      ----------
      query_units
      num_heads
      pos_embed_units
      max_distance
      bidirectional
      num_buckets
      method
      dropout
      dtype
      layout
      use_einsum
      """
      super().__init__()
      self._dropout = dropout
      self._method = method
      self._query_units = query_units
      self._num_heads = num_heads
      self._bidirectional = bidirectional
      self._num_buckets = num_buckets
      assert (
          query_units % num_heads == 0
      ), "The units must be divisible by the number of heads."
      self._head_query_units = query_units // num_heads
      self._max_distance = max_distance
      self._pos_embed_units = pos_embed_units
      self._dtype = dtype
      self._use_einsum = use_einsum
      self._layout = layout
      if self._layout not in ["NKT", "NTK", "TNK"]:
          raise ValueError(
              'layout="{}" is not supported'.format(self._layout)
          )
      if method == "transformer_xl":
          if pos_embed_units is None:
              pos_embed_units = self._num_heads * self._head_query_units
          self._rel_pos_embed = SinusoidalPositionalEmbedding(
              units=pos_embed_units, dtype=self._dtype
          )
          self._rel_proj = nn.Dense(
              units=query_units,
              in_units=pos_embed_units,
              flatten=False,
              use_bias=False,
              dtype=self._dtype,
          )
          self._dropout_layer = nn.Dropout(dropout)
      elif method == "shaw":
          assert (
              self._max_distance is not None
          ), 'Must set max_distance when method="shaw".'
          if self._bidirectional:
              vocab_size = self._max_distance * 2 + 1
          else:
              vocab_size = self._max_distance + 1
          self._rel_pos_embed = LearnedPositionalEmbedding(
              units=self._num_heads * self._head_query_units,
              max_length=vocab_size,
              weight_initializer=mx.init.Xavier(
                  rnd_type="gaussian", factor_type="in", magnitude=1
              ),
              mode="wrap" if self._bidirectional else "raise",
              dtype=self._dtype,
          )
      elif method == "t5":
          if self._num_buckets is None:
              self._num_buckets = 32
          if self._max_distance is None:
              self._max_distance = 128
          self._rel_pos_embed = BucketPositionalEmbedding(
              units=num_heads,
              num_buckets=self._num_buckets,
              max_distance=self._max_distance,
              bidirectional=self._bidirectional,
              embed_initializer=embed_initializer,
              dtype=self._dtype,
          )
      else:
          raise NotImplementedError(
              'method="{}" is currently not supported!'.format(method)
          )
  
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 1 to create a new RelAttentionScoreCell object from inherited base classes: HybridBlock with example: obj = RelAttentionScoreCell(query_units, num_heads, pos_embed_units, max_distance, bidirectional, num_buckets, method, dropout, dtype, layout, use_einsum, embed_initializer)