import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Weighted Sum implementation.
    Computes the weighted sum of values using corresponding weights.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, values: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
        """
        Compute weighted sum of values.

        Args:
            values (torch.Tensor): Input values [batch_size, feature_dim]
            weights (torch.Tensor): Corresponding weights [batch_size, feature_dim]

        Returns:
            torch.Tensor: Weighted sums [batch_size]
        """
        # 逐元素乘法
        elementwise_product = values * weights
        
        # 求和
        result = torch.sum(elementwise_product, dim=1)
        
        return result

batch_size = 1024
feature_dim = 512

def get_inputs():
    # Generate values and corresponding weights
    values = torch.randn(batch_size, feature_dim)
    weights = torch.rand(batch_size, feature_dim)  # Random weights between 0 and 1
    return [values, weights]

def get_init_inputs():
    return []  # No special initialization inputs needed
