LayerNorm

layernorm.py

class LayerNorm():
    def __init__(self, device, num_features):
        self.out = None
        self.gamma = torch.ones(num_features, device=device)
        self.bias = torch.zeros(num_features, device=device)

    def __call__(self, x: torch.Tensor):
        assert x.ndim == 2

        H = x.shape[1]
        avg = x.mean(dim=1, keepdim=True)
        std = torch.sqrt(1 / H * ((x - avg)**2).sum(dim=1, keepdim=True))

        self.out = (x - avg) / std * self.gamma + self.bias
        return self.out

    def params(self):
        return [self.gamma, self.bias]