def __init__(self, n_heads=1, dropout=0.0, mode='train'):
    """Returns a new :py:class:`PureAttention` instance.

    Args:
      n_heads: Number of attention heads.
      dropout: Probababilistic rate for attention dropout, which overrides
          (sets to zero) some attention strengths derived from query-key
          matching. As a result, on a given forward pass, some value vectors
          don't contribute to the output, analogous to how regular dropout can
          cause some node activations to be ignored. Applies only if layer is
          created in ``'train'`` mode.
      mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
    """
    super().__init__(n_in=4, n_out=2)
    self._n_heads = n_heads
    self._dropout = dropout
    self._mode = mode
Returns a new :py:class:`PureAttention` instance.

    Args:
      n_heads: Number of attention heads.
      dropout: Probababilistic rate for attention dropout, which overrides
          (sets to zero) some attention strengths derived from query-key
          matching. As a result, on a given forward pass, some value vectors
          don't contribute to the output, analogous to how regular dropout can
          cause some node activations to be ignored. Applies only if layer is
          created in ``'train'`` mode.
      mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
    
trax/layers/attention.py