"""FeedForward GELU (position-wise) — écrit à la main."""

from __future__ import annotations

import torch.nn as nn


class FeedForward(nn.Module):
    """MLP à deux couches : ``n_embd -> mult*n_embd -> n_embd`` avec GELU."""

    def __init__(self, n_embd: int, mult: int = 4, dropout: float = 0.0, bias: bool = False):
        super().__init__()
        hidden = mult * n_embd
        self.fc1 = nn.Linear(n_embd, hidden, bias=bias)
        self.fc2 = nn.Linear(hidden, n_embd, bias=bias)
        self.act = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.dropout(self.fc2(self.act(self.fc1(x))))