code: def _relative_position_to_absolute_position_masked(x):
    
    batch, heads, length, _ = common_layers.shape_list(x)
    x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [1, 0]])
    x = tf.reshape(x, [batch, heads, 1 + length, length])
    x = tf.slice(x, [0, 0, 1, 0], [-1, -1, -1, -1])
    return x
docs: Helper to dot_product_self_attention_relative_v2.
  
    Rearrange an attention logits or weights Tensor.
  
    The dimensions of the input represent:
    [batch, heads, query_position, memory_position - query_position + length - 1]
  
    The dimensions of the output represent:
    [batch, heads, query_position, memory_position]
  
    Only works with masked_attention.  Undefined behavior for regions of the
    input where memory_position > query_position.
  
    Args:
      x: a Tensor with shape [batch, heads, length, length]
  
    Returns:
      a Tensor with shape [batch, heads, length, length]