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


class SwiGLU(nn.Module):
    """Gated MLP: one path (with SiLU) gates another, then projects back down."""

    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        # gate and up both widen dim -> hidden_dim; down narrows back
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SiLU(gate) acts as a soft, per-element valve on the up path
        gated = F.silu(self.gate_proj(x)) * self.up_proj(x)
        return self.down_proj(gated)