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

class Model(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        torch.manual_seed(42)
        self.hidden_size = hidden_size
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.eps = 1e-5
    
    def forward(self, x):
        swish_out = x * torch.sigmoid(x)

        mean = swish_out.mean(-1, keepdim=True)
        var = swish_out.var(-1, keepdim=True, unbiased=False)
        swish_out = (swish_out - mean) / torch.sqrt(var + self.eps)

        output = swish_out * self.weight + self.bias
        
        return output

def get_inputs():
    batch_size = 4096
    seq_len = 128
    hidden_size = 768  # 典型的BERT hidden size
    x = torch.randn(batch_size, seq_len, hidden_size)
    return [x]

def get_init_inputs():
    return [768]
