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

class Model(nn.Module):
    """
    Chebyshev + LeakyReLU融合实现。
    先对输入应用LeakyReLU激活，然后计算Chebyshev距离。
    """
    def __init__(self, negative_slope=0.01):
        super(Model, self).__init__()
        self.negative_slope = negative_slope
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Compute Chebyshev + LeakyReLU fusion.

        Args:
            x (torch.Tensor): First set of vectors [batch_size, feature_dim]
            y (torch.Tensor): Second set of vectors [batch_size, feature_dim]

        Returns:
            torch.Tensor: Chebyshev distances after LeakyReLU activation [batch_size]
        """
        # Input validation
        if x.shape != y.shape:
            raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
        
        if x.dim() != 2:
            raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
        
        # Apply LeakyReLU activation to x
        x_activated = F.leaky_relu(x, negative_slope=self.negative_slope)
        
        # Compute Chebyshev distance: max(|x_activated - y|)
        chebyshev_dist = torch.max(torch.abs(x_activated - y), dim=1)[0]
        
        return chebyshev_dist

batch_size = 256
feature_dim = 512

def get_inputs():
    # Generate two sets of vectors
    x = torch.randn(batch_size, feature_dim)
    y = torch.randn(batch_size, feature_dim)
    return [x, y]

def get_init_inputs():
    return [0.01]  # negative_slope value
