import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class ModelNew(nn.Module):
    def __init__(self, in_features: int = 1024, out_features: int = 2048):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        fan_in = self.weight.size(1)
        bound = 1.0 / math.sqrt(fan_in)
        nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.linear(x, self.weight, self.bias)
        # 使用精确 GELU 以确保与基线一致的数值结果
        return F.gelu(y, approximate='none')


def get_init_inputs():
    return {"in_features": 1024, "out_features": 2048}


def get_inputs():
    B, T, D = 16, 512, 1024
    x = torch.randn(B, T, D)
    return x