import torch
import torch.nn as nn


class Model(nn.Module):
    """使用 PyTorch RMSNorm 的基准实现。"""

    def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
        super().__init__()
        if weight.dim() != 1:
            raise ValueError("RMSNorm 权重必须是一维向量。")
        feature_dim = weight.shape[0]
        self.rmsnorm = nn.RMSNorm(feature_dim, eps=eps, elementwise_affine=True)
        with torch.no_grad():
            self.rmsnorm.weight.copy_(weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """直接对输入做 RMSNorm，输出形状与输入一致。"""
        return self.rmsnorm(x)


batch_size = 16
feature_dim = 2048


def get_inputs():
    x = torch.randn(batch_size, feature_dim)
    return [x]


def get_init_inputs():
    weight = torch.randn(feature_dim)
    return [weight]