class_docstr: Relative attention layer.
  
  Layer that maps (location_bias, context_bias, pos_emb, q, k, v, mask)
      to (activations, mask).
  This layer type performs the inner workings of one pass of multi-head
  self-attention. It:
    - splits queries, keys, and values into multiple 'heads',
    - splits positional embeddings into multiple 'heads',
    - computes per-head attention weights from per-head (queries, keys),
    - applies mask to screen out positions that come from padding tokens,
    - [in `'train'` mode] applies dropout to attention weights,
    - uses attention weights to combine per-head values vectors, and
    - merges per-head results into outgoing activations matching original input
      activation vector shapes.
docstr: Returns a new PureAttention instance.
  
  Args:
    separate_cls: True/False if we separate_cls in calculations.
    n_heads: Number of attention heads.
    dropout: Probabilistic rate for dropout applied to attention strengths
        (based on query-key pairs) before applying them to values.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
docstr_tok: ["Returns","a","new","PureAttention","instance",".","Args",":","separate_cls",":","True\/False","if","we","separate_cls","in","calculations",".","n_heads",":","Number","of","attention","heads",".","dropout",":","Probabilistic","rate","for","dropout","applied","to","attention","strengths","(","based","on","query-key","pairs",")","before","applying","them","to","values",".","mode",":","One","of","`","'train","'","`",",","`","'eval","'","`",",","or","`","'predict","'","`","."]
code: def __init__(
      self, separate_cls, n_heads=1, dropout=0.0, mode="train"
  ):
      """Returns a new PureAttention instance.
  
      Args:
        separate_cls: True/False if we separate_cls in calculations.
        n_heads: Number of attention heads.
        dropout: Probabilistic rate for dropout applied to attention strengths
            (based on query-key pairs) before applying them to values.
        mode: One of `'train'`, `'eval'`, or `'predict'`.
      """
      super().__init__(n_in=7, n_out=2)
      self._separate_cls = separate_cls
      self._n_heads = n_heads
      self._dropout = dropout
      self._mode = mode
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 1 to create a new RelativeAttention object from inherited base classes: base with example: obj = RelativeAttention(separate_cls, n_heads, dropout, mode)