import torch
import torch.nn as nn

# --- 用于基准测试的配置 ---
BATCH_SIZE = 512
VECTORS = 8192
SHAPE = (BATCH_SIZE, VECTORS, 3)
DIM = -1

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.linalg.cross 作为基准模型。
    """
    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.cross(x, y, dim=self.dim)

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

def get_init_inputs():
    return [DIM]