import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, min_val, max_val):
        super(Model, self).__init__()
        self.min_val = min_val
        self.max_val = max_val
        self.range = max_val - min_val

    def forward(self, x):
        ranks = x.argsort(dim=-1).argsort(dim=-1).float()

        # Quantile Normalization to [0, 1]
        norm = ranks / (x.size(-1) - 1)

        # Range Expand
        return norm * self.range + self.min_val


batch_size = 256
dim = 256


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


def get_init_inputs():
    return [-1.0, 1.0]