import torch
import torch.nn as nn

BATCH_SIZE = 512
DIM = 4096
SHAPE = (BATCH_SIZE, DIM)
REDUCTION = 'mean'

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.nn.SoftMarginLoss 作为基准模型。
    """
    def __init__(self, reduction='mean'):
        super(Model, self).__init__()
        self.loss_fn = nn.SoftMarginLoss(reduction=reduction)
    
    def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
        return self.loss_fn(input_tensor, target_tensor)

def get_inputs():
    
    """
    生成用于测试的输入张量。
    """
    input_tensor = torch.randn(SHAPE, dtype=torch.float32)
    # target 张量必须只包含 1 和 -1
    # 使用 randint 生成 0 或 1，然后映射到 -1 或 1
    target_tensor = torch.randint(0, 2, SHAPE, dtype=torch.float32) * 2 - 1
    
    return [input_tensor.contiguous(), target_tensor.contiguous()]

def get_init_inputs():
    """
    提供模型初始化所需的参数。
    """
    return [REDUCTION]