code: 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,
      )
  
imports: import numpy
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
  from trax.layers.research import sparsity
label_desc: Use image node_id 6 for calling a global function with example usage: ConfigurableAttention(q_layer, k_layer, v_layer, final_layer, qkv_attention_layer, n_heads) and returns: cb