def Attention(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Returns a layer that maps `(vectors, mask)` to `(new_vectors, mask)`.

  This layer type represents one pass of multi-head self-attention, from vector
  set to vector set, using masks to represent out-of-bound (e.g., padding)
  positions. It:

    - makes three copies of incoming activations and maps these to multi-head
      query (Q) vectors, key (K) vectors, and value (V) vectors, respectively;
    - for each head, computes the scaled dot product of each Q-K pair;
    - applies mask to screen out positions that come from padding tokens
      (indicated by 0 value);
    - [in ``'train'`` mode] applies dropout to Q-K dot products;
    - for each head, computes Q-K attention strengths using a per-query softmax
      of the Q-K dot products;
    - for each head, for each query position, combines V vectors according
      to the Q-K attention strengths; and
    - concatenates and fuses resulting per-head vectors into outgoing
      activations matching original input activation shapes.

  Args:
    d_feature: Last/innermost dimension of activations in the input to and
        output from this layer.
    n_heads: Number of attention heads. Attention heads effectively split
        activation vectors into ``n_heads`` subvectors, of size
        ``d_feature / n_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'``.
  """
  return cb.Serial(
      cb.Select([0, 0, 0]),
      AttentionQKV(d_feature, n_heads=n_heads, dropout=dropout, mode=mode),
  )
Returns a layer that maps `(vectors, mask)` to `(new_vectors, mask)`.

  This layer type represents one pass of multi-head self-attention, from vector
  set to vector set, using masks to represent out-of-bound (e.g., padding)
  positions. It:

    - makes three copies of incoming activations and maps these to multi-head
      query (Q) vectors, key (K) vectors, and value (V) vectors, respectively;
    - for each head, computes the scaled dot product of each Q-K pair;
    - applies mask to screen out positions that come from padding tokens
      (indicated by 0 value);
    - [in ``'train'`` mode] applies dropout to Q-K dot products;
    - for each head, computes Q-K attention strengths using a per-query softmax
      of the Q-K dot products;
    - for each head, for each query position, combines V vectors according
      to the Q-K attention strengths; and
    - concatenates and fuses resulting per-head vectors into outgoing
      activations matching original input activation shapes.

  Args:
    d_feature: Last/innermost dimension of activations in the input to and
        output from this layer.
    n_heads: Number of attention heads. Attention heads effectively split
        activation vectors into ``n_heads`` subvectors, of size
        ``d_feature / n_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