| 10 | |
| 11 | |
| 12 | def PositionalEmbedding( |
| 13 | num_embeddings: int, |
| 14 | embedding_dim: int, |
| 15 | padding_idx: int, |
| 16 | learned: bool = False, |
| 17 | ): |
| 18 | if learned: |
| 19 | # if padding_idx is specified then offset the embedding ids by |
| 20 | # this index and adjust num_embeddings appropriately |
| 21 | # TODO: The right place for this offset would be inside |
| 22 | # LearnedPositionalEmbedding. Move this there for a cleaner implementation. |
| 23 | if padding_idx is not None: |
| 24 | num_embeddings = num_embeddings + padding_idx + 1 |
| 25 | m = LearnedPositionalEmbedding(num_embeddings, embedding_dim, padding_idx) |
| 26 | nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5) |
| 27 | if padding_idx is not None: |
| 28 | nn.init.constant_(m.weight[padding_idx], 0) |
| 29 | else: |
| 30 | m = SinusoidalPositionalEmbedding( |
| 31 | embedding_dim, |
| 32 | padding_idx, |
| 33 | init_size=num_embeddings + padding_idx + 1, |
| 34 | ) |
| 35 | return m |