import torch
import torch.nn as nn

BATCH_SIZE = 512
NUM_CLASSES = 4096 
REDUCTION = 'mean'
P = 1 # 1 for L1 hinge, 2 for L2
MARGIN = 1.0

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.nn.MultiMarginLoss 作为基准模型。
    """
    def __init__(self, p=1, margin=1.0, reduction='mean'):
        super(Model, self).__init__()
        # weight is not benchmarked for simplicity, but the CUDA kernel supports it.
        self.loss_fn = nn.MultiMarginLoss(p=p, margin=margin, 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(BATCH_SIZE, NUM_CLASSES, dtype=torch.float32)
    
    target_tensor = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    
    return [input_tensor.contiguous(), target_tensor.contiguous()]

def get_init_inputs():
    return [P, MARGIN, REDUCTION]