import torch
import torch.nn as nn

# --- 用于基准测试的配置 ---
BATCH_SIZE = 256
VECTORS = 2048
DIM_SIZE = 1024
SHAPE = (BATCH_SIZE, VECTORS, DIM_SIZE)
DIM = -1

# 使用 double 精度以保证数值稳定性
DTYPE = torch.float64

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.linalg.vecdot 作为基准模型。
    """
    def __init__(self, dim):
        super(Model, self).__init__()
        self.dim = dim
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return torch.linalg.vecdot(x, y, dim=self.dim)

def get_inputs():
    x = torch.randn(SHAPE, dtype=DTYPE)
    y = torch.randn(SHAPE, dtype=DTYPE)
    return [x.contiguous(), y.contiguous()]

def get_init_inputs():
    return [DIM]