(input, normalized_shape, eps=1e-6)
| 869 | |
| 870 | |
| 871 | def rms_norm(input, normalized_shape, eps=1e-6): |
| 872 | dtype = input.dtype |
| 873 | input = input.to(torch.float32) |
| 874 | variance = input.pow(2).flatten(-len(normalized_shape)).mean(-1)[(...,) + (None,) * len(normalized_shape)] |
| 875 | input = input * torch.rsqrt(variance + eps) |
| 876 | return input.to(dtype) |
| 877 | |
| 878 | class DiagonalGaussianDistribution(object): |
| 879 | def __init__(self, parameters, deterministic=False, rms_norm_mean=False, only_return_mean=False): |