| 163 | |
| 164 | |
| 165 | class RMSNorm(nn.Module): |
| 166 | def __init__(self, shape, eps=1e-6): |
| 167 | super().__init__() |
| 168 | self.eps = eps |
| 169 | self.scale = nn.Parameter(torch.ones(shape)) |
| 170 | |
| 171 | def extra_repr(self): |
| 172 | return f"shape={tuple(self.scale.shape)}, eps={self.eps}" |
| 173 | |
| 174 | def forward(self, x): |
| 175 | return rms_norm(x, self.scale, self.eps) |
| 176 | |
| 177 | |
| 178 | class AdaRMSNorm(nn.Module): |