code: def RelativeAttentionWrapper(
      d_feature,
      n_heads=1,
      dropout=0.0,
      max_inference_length=2048,
      mode="train",
      context_bias_layer=None,
      location_bias_layer=None,
      total_pooling=None,
  ):
      """Relative attention wrapper.
  
      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: dropout rate.
        max_inference_length: max inference length.
        mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
        context_bias_layer: context bias layer.
        location_bias_layer: location bias layer.
        total_pooling: total pooling.
  
      Returns:
        relative attention layer.
  
      Relative attention wrapper for compatibility with configurable attention,
      so that it can be called by `ApplyAttentionLayer`.
      """
      del max_inference_length
  
      attention = RelativeAttentionLMLayer(
          d_feature,
          context_bias_layer,
          location_bias_layer,
          total_pooling,
          n_heads=n_heads,
          dropout=dropout,
          mode=mode,
      )
  
      return cb.Serial(cb.Select([0, 0, 0]), attention)
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 1 for calling a global function with example usage: RelativeAttentionWrapper(d_feature, n_heads, dropout, max_inference_length, mode, context_bias_layer, location_bias_layer, total_pooling) and returns: cb