def get_rel_att_inputs(d_model, n_heads):
  """Global relative attentions bias initialization shared across layers."""
  assert d_model % n_heads == 0 and d_model % 2 == 0
  d_head = d_model // n_heads

  bias_initializer = init.RandomNormalInitializer(1e-6)
  context_bias_layer = core.Weights(
      bias_initializer, shape=(1, n_heads, 1, d_head))
  location_bias_layer = core.Weights(
      bias_initializer, shape=(1, n_heads, 1, d_head))
  return context_bias_layer, location_bias_layer
Global relative attentions bias initialization shared across layers.
trax/layers/research/rel_attention.py