text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import torch.nn as nn
  from .learned_positional_embedding import LearnedPositionalEmbedding, LearnedRelativePositionalEmbedding
  from .sinusoidal_positional_embedding import SinusoidalPositionalEmbedding
  
  
  def PositionalEmbedding(
          num_embeddings: int,
          embedding_dim: int,
          padding_idx: int,
          learned: bool = False,
  ):
      if learned:
          # if padding_idx is specified then offset the embedding ids by
          # this index and adjust num_embeddings appropriately
          # TODO: The right place for this offset would be inside
          # LearnedPositionalEmbedding. Move this there for a cleaner implementation.
          if padding_idx is not None:
              num_embeddings = num_embeddings + padding_idx + 1
          m = LearnedPositionalEmbedding(num_embeddings, embedding_dim, padding_idx)
          nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
          if padding_idx is not None:
              nn.init.constant_(m.weight[padding_idx], 0)
      else:
          m = SinusoidalPositionalEmbedding(
              embedding_dim, padding_idx, init_size=num_embeddings + padding_idx + 1,
          )
      return m
  
  
  def RelativePositionalEmbedding(
          max_relative_pos: int,
          num_heads: int,
          embedding_dim: int,
          unmasked: bool = False,
          heads_share_embeddings: bool = False,
          add_to_values: bool = False,
  ):
      return LearnedRelativePositionalEmbedding(
          max_relative_pos=max_relative_pos,
          num_heads=num_heads,
          embedding_dim=embedding_dim,
          heads_share_embeddings=heads_share_embeddings,
          unmasked=unmasked,
          add_to_values=add_to_values
      )
  