def ConfigurableAttention(q_layer, k_layer, v_layer, final_layer,  # pylint: disable=invalid-name
                          qkv_attention_layer, n_heads=1):
  """Returns a configured multi-head self-attention layer.

  A :py:class:`ConfigurableAttention` layer acts similarly to
  :py:class:`Attention` layers, but with configurable components. It

    - makes three copies of incoming activations and uses ``q_layer``,
      ``k_layer``, and ``v_layer`` to map activations to multi-head query (Q)
      vectors, key (K) vectors, and value (V) vectors, respectively;
    - uses ``qkv_attention_layer`` to compute per-head attention, similar to
      :py:class:`DotProductAttention` or :py:class:`DotProductCausalAttention`;
    - concatenates and fuses resulting per-head vectors into activations
      matching original input activation shapes; and
    - applies a final layer, ``final_layer``, mapping activations to
      activations (with shape matching the original input activations).

  Args:
    q_layer: Layer that maps input activations to per-head query activations.
    k_layer: Layer that maps input activations to per-head key activations.
    v_layer: Layer that maps input activations to per-head value activations.
    final_layer: After main multi-head computation and rejoining of heads,
        layer that maps activations to activations (with shape matching the
        original input activations).
    qkv_attention_layer: Layer the does the core multi-head self-attention
        computation.
    n_heads: Number of attention heads. Attention heads effectively split
        activation vectors into ``n_heads`` subvectors, of size
        ``d_feature / n_heads``.
  """
  return cb.Serial(
      cb.Branch(
          [q_layer, SplitIntoHeads(n_heads)],
          [k_layer, SplitIntoHeads(n_heads)],
          [v_layer, SplitIntoHeads(n_heads)],
      ),
      qkv_attention_layer,
      MergeHeads(n_heads),
      final_layer
  )
Returns a configured multi-head self-attention layer.

  A :py:class:`ConfigurableAttention` layer acts similarly to
  :py:class:`Attention` layers, but with configurable components. It

    - makes three copies of incoming activations and uses ``q_layer``,
      ``k_layer``, and ``v_layer`` to map activations to multi-head query (Q)
      vectors, key (K) vectors, and value (V) vectors, respectively;
    - uses ``qkv_attention_layer`` to compute per-head attention, similar to
      :py:class:`DotProductAttention` or :py:class:`DotProductCausalAttention`;
    - concatenates and fuses resulting per-head vectors into activations
      matching original input activation shapes; and
    - applies a final layer, ``final_layer``, mapping activations to
      activations (with shape matching the original input activations).

  Args:
    q_layer: Layer that maps input activations to per-head query activations.
    k_layer: Layer that maps input activations to per-head key activations.
    v_layer: Layer that maps input activations to per-head value activations.
    final_layer: After main multi-head computation and rejoining of heads,
        layer that maps activations to activations (with shape matching the
        original input activations).
    qkv_attention_layer: Layer the does the core multi-head self-attention
        computation.
    n_heads: Number of attention heads. Attention heads effectively split
        activation vectors into ``n_heads`` subvectors, of size
        ``d_feature / n_heads``.
  
trax/layers/attention.py