import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        """层归一化构造函数

        Args:
            dim (int): 层归一化的输入
            eps (float, optional): 偏置值. Defaults to 1e-6.
        """
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        
    def _norm(self, x):
        """归一化

        Args:
            x (_type_): 输入
        """
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim = True) + self.eps)
    
    def forward(self, x):
        """归一化层的前向传播

        Args:
            x (_type_): 输入
        Details:
            self.weight * self._norm：广播broadcast操作，输出的shape跟_norm后的shape一样，self.weight是可学习的
            type_as(x)：对归一化后的高精度 x 转为原来的类型 float16
        """
        return self.weight * self._norm(x.float()).type_as(x)

解释：