import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, scale, shift, eps=1e-5):
        super(Model, self).__init__()
        self.scale = scale
        self.shift = shift
        self.eps = eps

    def forward(self, x):
        min_x = x.min(dim=-1, keepdim=True).values
        max_x = x.max(dim=-1, keepdim=True).values

        range_x = max_x - min_x

        # MinMax Normalization
        norm = (x - min_x) / (range_x + self.eps)

        # Scale and Shift
        return norm * self.scale + self.shift


batch_size = 16
dim = 256


def get_inputs():
    x = torch.randn(batch_size, dim) * 10.0
    return [x]


def get_init_inputs():
    return [2.0, 1.0]