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
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'`.
    
trax/layers/research/rel_attention.py