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

class Model(nn.Module):
    """
    Cosine Distance implementation - 标准版本
    """
    def __init__(self, eps=1e-8):
        super(Model, self).__init__()
        self.eps = eps
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Compute cosine distance between x and y.
        Cosine Distance = 1 - Cosine Similarity
        
        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: Cosine distances [batch_size], range [0, 2]
        """
        # 使用PyTorch内置的cosine_similarity函数
        cosine_sim = F.cosine_similarity(x, y, dim=1, eps=self.eps)
        
        # 转换为距离
        cosine_dist = 1.0 - cosine_sim
        
        # 确保结果在合理范围内
        return torch.clamp(cosine_dist, 0.0, 2.0)

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 [1e-8]  # eps value
