import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    """
    BCE with Sigmoid implementation for binary classification.
    使用binary_cross_entropy_with_logits作为基准
    """
    def __init__(self, reduction='sum'):
        super(Model, self).__init__()
        self.reduction = reduction
    
    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """
        Compute BCE Loss with logits.
        
        Args:
            inputs (torch.Tensor): Predicted logits
            targets (torch.Tensor): Ground truth labels (0 or 1)
        
        Returns:
            torch.Tensor: Computed BCE loss
        """
        # 确保类型一致
        inputs = inputs.to(torch.float32)
        targets = targets.to(torch.float32)
        
        # 使用PyTorch的标准实现
        return F.binary_cross_entropy_with_logits(
            inputs, 
            targets, 
            reduction=self.reduction
        )

batch_size = 256
num_features = 2000

def get_inputs():
    # 生成logits（不是概率）
    input_logits = torch.randn(batch_size, num_features, dtype=torch.float32)
    target_labels = torch.randint(0, 2, (batch_size, num_features), dtype=torch.float32)
    return [input_logits, target_labels]

def get_init_inputs():
    return []
