code: def _matmul_with_relative_keys_2d(x, y, heads_share_relative_embedding):
    
    if heads_share_relative_embedding:
      ret = tf.einsum("bhxyd,md->bhxym", x, y)
    else:
      ret = tf.einsum("bhxyd,hmd->bhxym", x, y)
    return ret
docs: Helper function for dot_product_unmasked_self_attention_relative_2d.