| 64 | class RmsNorm(Module): |
| 65 | |
| 66 | def __init__(self, |
| 67 | normalized_shape, |
| 68 | num_groups=1, |
| 69 | eps=1e-06, |
| 70 | elementwise_affine=True, |
| 71 | dtype=None): |
| 72 | super().__init__() |
| 73 | if isinstance(normalized_shape, int): |
| 74 | normalized_shape = (normalized_shape, ) |
| 75 | self.normalized_shape = tuple(normalized_shape) |
| 76 | self.elementwise_affine = elementwise_affine |
| 77 | self.num_groups = num_groups |
| 78 | num_channels = normalized_shape[-1] |
| 79 | if num_channels % num_groups != 0: |
| 80 | raise ValueError('num_channels must be divisible by num_groups') |
| 81 | if self.elementwise_affine: |
| 82 | self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) |
| 83 | else: |
| 84 | self.register_parameter('weight', None) |
| 85 | |
| 86 | self.eps = eps |
| 87 | self.dtype = dtype |
| 88 | |
| 89 | def forward(self, x, normalized_shape=None): |
| 90 | weight = None if self.weight is None else self.weight.value |