| 578 | return x + skip |
| 579 | |
| 580 | class MappingNetwork(nn.Module): |
| 581 | def __init__(self, n_layers, d_model, d_ff, dropout=0.0): |
| 582 | super().__init__() |
| 583 | self.in_norm = RMSNorm(d_model) |
| 584 | self.blocks = nn.ModuleList([MappingFeedForwardBlock(d_model, d_ff, dropout=dropout) for _ in range(n_layers)]) |
| 585 | self.out_norm = RMSNorm(d_model) |
| 586 | |
| 587 | def forward(self, x): |
| 588 | x = self.in_norm(x) |
| 589 | for block in self.blocks: |
| 590 | x = block(x) |
| 591 | x = self.out_norm(x) |
| 592 | return x |
| 593 | |
| 594 | |
| 595 | # Token merging and splitting |