import torch
import torch.nn as nn

class SwishTorchModel(nn.Module):
    """PyTorch原生Swish激活函数实现"""
    
    def __init__(self):
        super(SwishTorchModel, self).__init__()
    
    def forward(self, x):
        return x * torch.sigmoid(x)

def get_init_inputs():
    """获取模型初始化参数"""
    return []

def get_inputs():
    """获取模型输入数据"""
    torch.manual_seed(42)
    return [torch.randn(128, 256)]

def test_swish():
    """测试Swish激活函数"""
    print("=" * 60)
    print("Swish激活函数测试")
    print("=" * 60)
    
    # 创建测试数据
    torch.manual_seed(42)
    input_tensor = torch.randn(128, 256)
    
    # 创建模型
    model = SwishTorchModel()
    
    # 前向传播
    output = model(input_tensor)
    
    print(f"输入形状: {input_tensor.shape}")
    print(f"输出形状: {output.shape}")
    print(f"输入范围: [{input_tensor.min().item():.3f}, {input_tensor.max().item():.3f}]")
    print(f"输出范围: [{output.min().item():.3f}, {output.max().item():.3f}]")
    
    # 验证Swish函数
    expected_output = input_tensor * torch.sigmoid(input_tensor)
    diff = torch.abs(output - expected_output).max().item()
    print(f"与预期输出最大差异: {diff:.6f}")
    
    if diff < 1e-6:
        print("✅ Swish激活函数实现正确")
    else:
        print("❌ Swish激活函数实现有误")
    
    return output

def main():
    """主测试函数"""
    try:
        output = test_swish()
        print("\n" + "=" * 60)
        print("测试总结")
        print("=" * 60)
        print("✅ Swish激活函数测试完成")
        
    except Exception as e:
        print(f"测试过程中出现错误: {e}")

if __name__ == "__main__":
    main()