import torch
import torch.nn as nn

BATCH_SIZE = 2048
SEQ_LEN = 2048  
SHAPE = (BATCH_SIZE, SEQ_LEN)
DIM = 1

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

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

def get_init_inputs():
    return [DIM]