code: def RelativeAttentionLMLayer(
      d_feature,
      context_bias_layer,
      location_bias_layer,
      total_kv_pooling,
      separate_cls=False,
      n_heads=1,
      dropout=0.0,
      mode="train",
  ):
      """Returns a layer that maps (q, k, v) to (activations).
  
      Same as standard Relative attention layer but additionally based on sizes
      of queries and keys prepares a mask that masks out the future.
      Masking the future is the concept primarily used for Language Modelling.
      Args:
        d_feature: Depth/dimensionality of feature embedding.
        context_bias_layer: Global context bias from Transformer XL's attention.
          There should be one such layer shared for all relative attention layers
        location_bias_layer: Global location bias from Transformer XL's attention.
          There should be one such layer shared for all relative attention layers.
        total_kv_pooling: Accumulated pool size of keys/values used at this layer.
        separate_cls: True/False if we separate_cls in calculations.
        n_heads: Number of attention heads.
        dropout: Probabilistic rate for internal dropout applied to attention
          activations (based on query-key pairs) before dotting them with values.
        mode: One of `'train'`, `'eval'`, or `'predict'`.
      """
  
      attention = RelativeAttentionLayer(
          d_feature,
          context_bias_layer,
          location_bias_layer,
          total_kv_pooling,
          separate_cls,
          n_heads=n_heads,
          dropout=dropout,
          mode=mode,
      )
  
      return cb.Serial(
          CreateAttentionMaskLayer(),  # q, k, v, mask
          attention,  # vecs, mask
          cb.Select([0], n_in=2),  # vecs
      )
  
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 4 for calling a global function with example usage: RelativeAttentionLMLayer(d_feature, context_bias_layer, location_bias_layer, total_kv_pooling, separate_cls, n_heads, dropout, mode) and returns: cb