| 49 | |
| 50 | |
| 51 | class MLP(nn.Module): |
| 52 | def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int, init_scale: float): |
| 53 | super().__init__() |
| 54 | self.width = width |
| 55 | self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype) |
| 56 | self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype) |
| 57 | self.gelu = nn.GELU() |
| 58 | init_linear(self.c_fc, init_scale) |
| 59 | init_linear(self.c_proj, init_scale) |
| 60 | |
| 61 | def forward(self, x): |
| 62 | return self.c_proj(self.gelu(self.c_fc(x))) |
| 63 | |
| 64 | |
| 65 | class QKVMultiheadAttention(nn.Module): |