import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Huber Loss implementation - robust loss function that is quadratic for small errors and linear for large errors.
    """
    def __init__(self, delta=1.0):
        super(Model, self).__init__()
        self.delta = delta
    
    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Compute Huber loss between input and target tensors.

        Args:
            input (torch.Tensor): Predicted values [batch_size, ...]
            target (torch.Tensor): Target values [batch_size, ...]

        Returns:
            torch.Tensor: Scalar Huber loss value
        """
        # 计算绝对误差
        abs_error = torch.abs(input - target)
        
        # 创建二次区域和线性区域的mask
        quadratic_mask = abs_error <= self.delta
        linear_mask = ~quadratic_mask
        
        # 二次区域：0.5 * error²
        quadratic_loss = 0.5 * torch.pow(input - target, 2)
        
        # 线性区域：delta * (|error| - 0.5 * delta)
        linear_loss = self.delta * (abs_error - 0.5 * self.delta)
        
        # 组合损失
        loss = torch.where(quadratic_mask, quadratic_loss, linear_loss)
        
        return torch.sum(loss)

batch_size = 512
feature_dim = 512
delta = 1.0

def get_inputs():
    input = torch.randn(batch_size, feature_dim)
    target = torch.randn(batch_size, feature_dim)
    return [input, target]

def get_init_inputs():
    return [delta]  # delta parameter
