import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, tl=-1.0, al=0.1, tr=1.0, ar=0.1):
        super().__init__()
        self.tl = tl
        self.al = al
        self.tr = tr
        self.ar = ar

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y_left = self.tl + self.al * (x - self.tl)

        y_right = self.tr + self.ar * (x - self.tr)

        y_mid_and_right = torch.where(x < self.tr, x, y_right)

        return torch.where(x <= self.tl, y_left, y_mid_and_right)


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, 0.1, 1.0, 0.1]