import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs LayerNorm normalization using PyTorch's built-in nn.LayerNorm.
    """

    def __init__(self, normalized_shape=None, eps=1e-5, elementwise_affine=True):
        super(Model, self).__init__()
        # 如果未指定normalized_shape，将在forward中动态设置
        self.normalized_shape = normalized_shape
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        self.layernorm = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies LayerNorm to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of any shape.

        Returns:
            torch.Tensor: Output tensor with LayerNorm applied, same shape as input.
        """
        # 如果layernorm未初始化，根据输入形状动态创建
        if self.layernorm is None:
            if self.normalized_shape is None:
                # 默认对最后一个维度进行归一化
                self.normalized_shape = x.shape[1:]
            self.layernorm = nn.LayerNorm(
                normalized_shape=self.normalized_shape,
                eps=self.eps,
                elementwise_affine=self.elementwise_affine
            ).to(x.device)

        return self.layernorm(x)


batch_size = 16
dim = 16384


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


def get_init_inputs():
    # 可以传入归一化形状、eps等参数，保持向后兼容
    return []  # 使用默认参数