| 536 | class LayerNorm(nn.Module): |
| 537 | |
| 538 | def __init__( |
| 539 | self, |
| 540 | hidden_size: int, |
| 541 | elementwise_affine: bool = True, |
| 542 | bias: bool = False, |
| 543 | eps: float = 1e-5 |
| 544 | ) -> LayerNorm: |
| 545 | super().__init__() |
| 546 | |
| 547 | self.hidden_size = hidden_size |
| 548 | self.elementwise_affine = elementwise_affine |
| 549 | self.eps = eps |
| 550 | |
| 551 | self.register_parameter("weight", None) |
| 552 | self.register_parameter("bias", None) |
| 553 | if elementwise_affine: |
| 554 | self.weight = nn.Parameter(torch.ones(hidden_size)) |
| 555 | if bias: |
| 556 | self.bias = nn.Parameter(torch.zeros(hidden_size)) |
| 557 | |
| 558 | def __repr__(self) -> str: |
| 559 | s = f"{self.__class__.__name__}({self.hidden_size}" |