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)
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`.
  
trax/layers/research/rel_attention.py