| 328 | |
| 329 | |
| 330 | class T5LayerNorm(torch.nn.Module): |
| 331 | def __init__(self, hidden_size, eps=1e-6, dtype=None, device=None): |
| 332 | super().__init__() |
| 333 | self.weight = torch.nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) |
| 334 | self.variance_epsilon = eps |
| 335 | |
| 336 | def forward(self, x): |
| 337 | variance = x.pow(2).mean(-1, keepdim=True) |
| 338 | x = x * torch.rsqrt(variance + self.variance_epsilon) |
| 339 | return self.weight.to(device=x.device, dtype=x.dtype) * x |
| 340 | |
| 341 | |
| 342 | class T5DenseGatedActDense(torch.nn.Module): |