import torch
import torch.nn as nn

BATCH_SIZE = 4096
FEATURES = 2048
SHAPE = (BATCH_SIZE, FEATURES)
DIM = 1

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.logsumexp 作为基准模型。
    """
    def __init__(self, dim, keepdim=False):
        super(Model, self).__init__()
        self.dim = dim
        self.keepdim = keepdim
    
    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return torch.logsumexp(input_tensor, dim=self.dim, keepdim=self.keepdim)

def get_inputs():
    # 扩大值的范围以测试数值稳定性
    input_tensor = torch.randn(SHAPE, dtype=torch.float32) * 10.0
    return [input_tensor.contiguous()]

def get_init_inputs():
    return [DIM, False] # keepdim=False