| 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) |