| 492 | |
| 493 | |
| 494 | class FeedForwardBlock(nn.Module): |
| 495 | def __init__(self, d_model, d_ff, cond_features, dropout=0.0): |
| 496 | super().__init__() |
| 497 | self.norm = AdaRMSNorm(d_model, cond_features) |
| 498 | self.up_proj = apply_wd(LinearGEGLU(d_model, d_ff, bias=False)) |
| 499 | self.dropout = nn.Dropout(dropout) |
| 500 | self.down_proj = apply_wd(zero_init(Linear(d_ff, d_model, bias=False))) |
| 501 | |
| 502 | def forward(self, x, cond): |
| 503 | skip = x |
| 504 | x = self.norm(x, cond) |
| 505 | x = self.up_proj(x) |
| 506 | x = self.dropout(x) |
| 507 | x = self.down_proj(x) |
| 508 | return x + skip |
| 509 | |
| 510 | class GlobalTransformerLayer(nn.Module): |
| 511 | def __init__(self, d_model, d_ff, d_head, cond_features, dropout=0.0): |