| 24 | |
| 25 | |
| 26 | class LayerNorm(Module): |
| 27 | |
| 28 | def __init__(self, |
| 29 | normalized_shape, |
| 30 | eps=1e-05, |
| 31 | elementwise_affine=True, |
| 32 | bias=True, |
| 33 | dtype=None, |
| 34 | tp_size=1, |
| 35 | tp_dim=-1): |
| 36 | super().__init__() |
| 37 | if isinstance(normalized_shape, int): |
| 38 | normalized_shape = (normalized_shape, ) |
| 39 | self.normalized_shape = tuple(normalized_shape) |
| 40 | self.elementwise_affine = elementwise_affine |
| 41 | if self.elementwise_affine: |
| 42 | self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) |
| 43 | if bias: |
| 44 | self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) |
| 45 | else: |
| 46 | self.register_parameter('bias', None) |
| 47 | else: |
| 48 | self.register_parameter('weight', None) |
| 49 | self.register_parameter('bias', None) |
| 50 | |
| 51 | self.eps = eps |
| 52 | self.dtype = dtype |
| 53 | self.tp_size = tp_size |
| 54 | self.tp_dim = tp_dim |
| 55 | |
| 56 | def forward(self, x, normalized_shape=None): |
| 57 | weight = 1. if self.weight is None else self.weight.value |
| 58 | bias = 0. if self.bias is None else self.bias.value |
| 59 | if normalized_shape is None: |
| 60 | normalized_shape = self.normalized_shape |
| 61 | return layer_norm(x, normalized_shape, weight, bias, self.eps) |
| 62 | |
| 63 | |
| 64 | class RmsNorm(Module): |
no outgoing calls
no test coverage detected