"""Embedding layers used by TaoNet."""

import torch.nn as nn


class FactorizedEmbedding(nn.Module):
    """Low-rank factorized embedding."""

    def __init__(self, vocab_size, d_model, d_embed_rank=96):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.d_embed_rank = d_embed_rank

        self.embed = nn.Embedding(vocab_size, d_embed_rank)
        self.proj = nn.Linear(d_embed_rank, d_model, bias=False)

        nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.proj.weight, mean=0.0, std=0.02)

    def forward(self, input_ids):
        x = self.embed(input_ids)
        return self.proj(x)