def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train',
                 cache_KV_in_predict=False, q_sparsity=None,
                 result_sparsity=None):
  """Returns a layer that maps `(AQ, AK, AV, mask)` to `(new-A, mask)`.

  Unlike :py:class:`Attention` above, :py:class:`AttentionQKV` allows the
  incoming activations (`AQ`, `AK`, and `AV`) to come from different sources.
  This is used, for instance, in encoder-decoder attention (Q-related
  activations `AQ` from the decoder, K- and V-related activations -- `AK` and
  `AV` -- from the encoder). Otherwise, see the :py:class:`Attention`
  description for further context/details.

  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'``.
    cache_KV_in_predict: Whether to cache K/V arrays in ``'predict'`` mode.
    q_sparsity: Sparsity with which to process queries. If ``None``,
        :py:class:`Dense` is used; if ``'noop'``, no processing is used.
    result_sparsity: Sparsity with which to process result of the attention.
        If ``None``, :py:class:`Dense` is used; if ``'noop'``, no processing is
        used.
  """
  def _SparsifiableDense(layer_sparsity):
    if layer_sparsity is None:
      return core.Dense(d_feature)
    elif layer_sparsity == 'noop':
      return cb.Serial()  # No-op layer.
    else:
      d_module = d_feature // layer_sparsity
      return cb.Serial(
          sparsity.FactoredDense(layer_sparsity, d_feature, d_feature),
          sparsity.LocallyConvDense(layer_sparsity, d_module, mode=mode,
                                    kernel_size=3, length_kernel_size=3)
      )

  def _CacheableDense():
    if cache_KV_in_predict and mode == 'predict':
      return cb.Cache(core.Dense(d_feature))
    else:
      return core.Dense(d_feature)

  def _PureAttention():
    return PureAttention(n_heads=n_heads, dropout=dropout, mode=mode)

  return cb.Serial(
      cb.Parallel(_SparsifiableDense(q_sparsity),
                  _CacheableDense(),
                  _CacheableDense()),
      _PureAttention(),
      _SparsifiableDense(result_sparsity),
  )
Returns a layer that maps `(AQ, AK, AV, mask)` to `(new-A, mask)`.

  Unlike :py:class:`Attention` above, :py:class:`AttentionQKV` allows the
  incoming activations (`AQ`, `AK`, and `AV`) to come from different sources.
  This is used, for instance, in encoder-decoder attention (Q-related
  activations `AQ` from the decoder, K- and V-related activations -- `AK` and
  `AV` -- from the encoder). Otherwise, see the :py:class:`Attention`
  description for further context/details.

  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'``.
    cache_KV_in_predict: Whether to cache K/V arrays in ``'predict'`` mode.
    q_sparsity: Sparsity with which to process queries. If ``None``,
        :py:class:`Dense` is used; if ``'noop'``, no processing is used.
    result_sparsity: Sparsity with which to process result of the attention.
        If ``None``, :py:class:`Dense` is used; if ``'noop'``, no processing is
        used.
  
trax/layers/attention.py