import torch
import torch.nn as nn

class LayerNormal(nn.Module):
    def __init__(self,emb_dim):
        super().__init__()
        # 添加一个小浮点数，避免一组相同的常数造成方差为0（方差>=0）
        self.eps = 1e-5
        # 偏移和缩放权重参数
        self.scale = nn.Parameter(torch.ones(emb_dim))
        # print(f"scale.shape:{self.scale.shape}")
        self.shift = nn.Parameter(torch.zeros(emb_dim))
        # print(f"shift.shape:{self.shift.shape}")
    
    def forward(self,x):
        mean = x.mean(dim=-1,keepdim=True)
        # 为什么这里需要unbiased？
        """ 
        种设置确保我们的方法与第 6 章中将加载的预训练权重兼容。
         """
        var = x.var(dim=-1,keepdim=True,unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        # print(f"norm_x.shape:{norm_x.shape}")
        return self.scale * norm_x + self.shift