"""
SwiGLU feed-forward network (Llama-style).

Two parallel projections up to the hidden width: one is the 'gate', run
through SiLU; the other is the 'value'. Their elementwise product is
projected back down. SiLU(gate) * up acts like a learned, smooth filter.
"""

import torch.nn as nn
import torch.nn.functional as F


class AbyssSwiGLU(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate = nn.Linear(config.n_embd, config.n_inner, bias=config.bias)
        self.up = nn.Linear(config.n_embd, config.n_inner, bias=config.bias)
        self.down = nn.Linear(config.n_inner, config.n_embd, bias=config.bias)
        self.drop = nn.Dropout(config.dropout)

    
    def forward(self, x):
        return self.drop(self.down(F.silu(self.gate(x)) * self.up(x)))