import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, dim=-1):
        super(Model, self).__init__()
        self.dim = dim
    
    def forward(self, x):
        # 简化实现：直接使用PyTorch内置Softmax
        return torch.softmax(x, dim=self.dim)

def get_inputs():
    """生成适合Softmax测试的输入数据"""
    torch.manual_seed(42)
    
    # 使用更适合CUDA优化的数据规模
    batch_size = 128  # 减少批量大小，避免内存限制
    seq_length = 512   # 减少序列长度，更适合CUDA优化
    
    # 生成简单的测试数据
    input_data = torch.randn(batch_size, seq_length) * 0.1
    
    return [input_data]

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

def generate_test_cases():
    """生成多种测试用例，用于全面测试Softmax性能"""
    test_cases = []
    
    # 小规模测试用例 - 测试warp级优化
    test_cases.append({
        'name': 'small_batch_small_seq',
        'input': torch.randn(64, 128) * 0.1,  # 适合warp级处理
        'description': '小批量小序列测试（warp优化）'
    })
    
    # 中等规模测试用例 - 测试块级优化
    test_cases.append({
        'name': 'medium_batch_medium_seq',
        'input': torch.randn(256, 512) * 0.15,  # 适合块级处理
        'description': '中等批量中等序列测试（块优化）'
    })
    
    # 大规模测试用例 - 测试分层归约优化
    test_cases.append({
        'name': 'large_batch_large_seq',
        'input': torch.randn(512, 2048) * 0.2,  # 适合分层归约
        'description': '大批量大序列测试（分层归约）'
    })
    
    # 超大规模测试用例 - 测试多遍归约优化
    test_cases.append({
        'name': 'huge_batch_huge_seq',
        'input': torch.randn(1024, 4096) * 0.25,  # 适合多遍归约
        'description': '超大批量超长序列测试（多遍归约）'
    })
    
    # 极端情况测试用例 - 测试数值稳定性
    test_cases.append({
        'name': 'extreme_values',
        'input': torch.tensor([[100.0, -100.0, 50.0, -50.0, 200.0, -200.0]]),
        'description': '极端数值稳定性测试'
    })
    
    # 边界情况测试用例 - 测试边界处理
    test_cases.append({
        'name': 'boundary_cases',
        'input': torch.tensor([[1e-10, 1e10, 0.0, -1e10, 1e-5]]),
        'description': '边界数值处理测试'
    })
    
    return test_cases