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

# --- 用于基准测试的配置 ---
BATCH_SIZE = 256
CHANNELS = 1024
HEIGHT = 32
WIDTH = 32
SHAPE = (BATCH_SIZE, CHANNELS, HEIGHT, WIDTH)

# GeM Pooling 的超参数
P_VALUE = 3.0
EPS_VALUE = 1e-6

class GemPooling(nn.Module):
    """
    待优化的算子类：Generalized Mean Pooling (GeM)
    公式: f(X) = (mean(clamp(X, min=eps)^p))^(1/p)
    """
    def __init__(self, p=3.0, eps=1e-6):
        super(GemPooling, self).__init__()
        self.p = p
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.clamp(min=self.eps)
        x = x.pow(self.p)
        # 输入形状 (N, C, H, W) -> 输出形状 (N, C)
        x = x.mean(dim=[-2, -1])
        x = x.pow(1.0 / self.p)
        
        return x

class Model(nn.Module):
    def __init__(self, p=3.0, eps=1e-6):
        super(Model, self).__init__()
        self.gem_pooling = GemPooling(p=p, eps=eps)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.gem_pooling(x)

def get_inputs():
    input_tensor = torch.randn(SHAPE, dtype=torch.float32)
    return [input_tensor.contiguous()]

def get_init_inputs():
    return [P_VALUE, EPS_VALUE]