| 271 | |
| 272 | |
| 273 | class RMSNorm(Module): |
| 274 | def __init__(self, dim, channel_first=False, images=False, bias=False): |
| 275 | super().__init__() |
| 276 | broadcastable_dims = (1, 1, 1) if not images else (1, 1) |
| 277 | shape = (dim, *broadcastable_dims) if channel_first else (dim,) |
| 278 | |
| 279 | self.channel_first = channel_first |
| 280 | self.scale = dim**0.5 |
| 281 | self.gamma = nn.Parameter(torch.ones(shape)) |
| 282 | self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 |
| 283 | |
| 284 | def forward(self, x): |
| 285 | return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias |
| 286 | |
| 287 | |
| 288 | class AdaptiveRMSNorm(Module): |