(self, x: torch.Tensor)
| 832 | self.dim = dim |
| 833 | |
| 834 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 835 | # NOTE: the original RMSNorm paper implementation is not equivalent |
| 836 | norm_x = torch.mean(x * x, dim=self.dim, keepdim=True) |
| 837 | x_normed = x * torch.rsqrt(norm_x + self.eps) |
| 838 | return self.weight * x_normed |
| 839 | |
| 840 | def reset_parameters(self): |
| 841 | torch.nn.init.ones_(self.weight) |
nothing calls this directly
no outgoing calls
no test coverage detected