# localresponsenorm_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F


N, C, H, W = 16, 64, 32, 32

SIZE = 5      # 窗口大小 (n)
ALPHA = 1e-4  # 缩放因子 (alpha)
BETA = 0.75   # 幂指数 (beta)
K = 2.0       # 偏置项 (k)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.lrn = nn.LocalResponseNorm(SIZE, alpha=ALPHA, beta=BETA, k=K)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.lrn(x)

def get_inputs():
    x = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x]

def get_init_inputs():
    return []