"""
基础层模块
包含Linear和LayerNorm层
"""

import numpy as np


class Linear:
    """全连接层（线性层）"""
    def __init__(self, in_features, out_features):
        # 使用He初始化权重
        self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features)
        # 偏置初始化为0
        self.bias = np.zeros(out_features)
        
    def forward(self, x):
        """前向传播：y = xW + b"""
        return x @ self.weight + self.bias


class LayerNorm:
    """层归一化（Layer Normalization）"""
    def __init__(self, d_model, eps=1e-6):
        # 缩放参数
        self.gamma = np.ones(d_model)
        # 偏移参数
        self.beta = np.zeros(d_model)
        # 防止除零的小常数
        self.eps = eps
        
    def forward(self, x):
        """前向传播：对最后一个维度进行归一化"""
        mean = np.mean(x, axis=-1, keepdims=True)
        std = np.std(x, axis=-1, keepdims=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta