import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Cosine Similarity implementation.
    Computes the cosine similarity between two sets of vectors.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Compute cosine similarity between x and y.

        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 similarities [batch_size]
        """
        # 计算L2范数
        norm_x = torch.sqrt(torch.sum(x * x, dim=1, keepdim=True))
        norm_y = torch.sqrt(torch.sum(y * y, dim=1, keepdim=True))
        
        # 计算点积
        dot_product = torch.sum(x * y, dim=1, keepdim=True)
        
        # 计算余弦相似度，避免除零
        cosine_sim = dot_product / (norm_x * norm_y + 1e-8)
        
        return cosine_sim.squeeze(1)

batch_size = 1024
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 []  # No special initialization inputs needed
