import torch
import torch.nn as nn

class FeedForward(nn.Module):
    def __init__(self,config):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(config['emb_dim'], 4 * config['emb_dim']),
            GELU(),
            nn.Linear(4 * config['emb_dim'], config['emb_dim']),
        )
    
    def forward(self,x):
        return self.layers(x)

class GELU(nn.Module):
    def __init__(self,):
        super().__init__()

    def forward(self,x):
        return 0.5 * x * (
            1 + torch.tanh(
                torch.sqrt(
                    torch.tensor(2.0 / torch.pi, device=x.device)
                ) * (x + 0.044715 * (x**3))
            )
        )