import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, alpha=1.0, c=1.0):
        super().__init__()
        self.alpha = alpha
        self.c = c

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # PLU Formula: max(alpha(x+c) - c, min(alpha(x-c) + c, x))
        term1 = self.alpha * (x + self.c) - self.c
        term2 = self.alpha * (x - self.c) + self.c

        inner_min = torch.min(term2, x)

        return torch.max(term1, inner_min)


batch_size = 128
feature_dim = 512


def get_inputs():
    x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [1.0, 1.0]