import torch
import torch.nn as nn
import numpy as np

BATCH_SIZE = 512
NUM_CLASSES = 1024
REDUCTION = 'mean'
# 每个样本的正类标签数量范围
MIN_LABELS = 1
MAX_LABELS = 10

class Model(nn.Module):
    def __init__(self, reduction='mean'):
        super(Model, self).__init__()
        self.loss_fn = nn.MultiLabelMarginLoss(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每行包含正类索引，并用 -1 填充
    target_np = np.full((BATCH_SIZE, NUM_CLASSES), -1, dtype=np.int64)
    for i in range(BATCH_SIZE):
        num_labels = np.random.randint(MIN_LABELS, MAX_LABELS + 1)
        labels = np.random.choice(NUM_CLASSES, num_labels, replace=False)
        target_np[i, :num_labels] = labels
    target_tensor = torch.from_numpy(target_np)

    return [input_tensor.contiguous(), target_tensor.contiguous()]

def get_init_inputs():
    return [REDUCTION]