| 86 | |
| 87 | |
| 88 | class SharedEmbedding(nn.Embedding): |
| 89 | def forward(self, input: torch.Tensor, unembed: bool = False) -> torch.Tensor: |
| 90 | if unembed: |
| 91 | return F.linear(input, self.weight) |
| 92 | return super().forward(input) |
| 93 | |
| 94 | |
| 95 | class MPTAttentionFused(nn.Module): |