import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, in_features, alpha=0.0):
        super().__init__()
        self.alpha = nn.Parameter(torch.full((in_features,), alpha))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shape = [1] * x.dim()
        shape[1] = -1
        a = self.alpha.view(shape)

        condition_zero = (a == 0)
        condition_pos = (a > 0)
        condition_neg = (a < 0)

        res = torch.zeros_like(x)

        if condition_zero.any():
            res = torch.where(condition_zero, x, res)

        if condition_pos.any():
            res = torch.where(condition_pos, (torch.exp(a * x) - 1.0) / a + a, res)

        if condition_neg.any():
            res = torch.where(condition_neg, -torch.log(1.0 - a * (x + a)) / a, res)

        return res


batch_size = 128
in_features = 64
height = 64
width = 64


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


def get_init_inputs():
    return [in_features, 0.5]