import torch
import torch.nn as nn

BATCH_SIZE = 2048
DIM = 2048
SHAPE = (BATCH_SIZE, DIM)

QMIN = -128
QMAX = 127

class FakeQuantize(nn.Module):
    """
    Fake Quantize Per Tensor Affine
    """
    def __init__(self, quant_min=QMIN, quant_max=QMAX):
        super(FakeQuantize, self).__init__()
        self.quant_min = quant_min
        self.quant_max = quant_max

    def forward(self, x, scale, zero_point):
        return torch.fake_quantize_per_tensor_affine(
            x, scale, zero_point, self.quant_min, self.quant_max
        )

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fq = FakeQuantize()
    
    def forward(self, x, scale, zero_point):
        return self.fq(x, scale, zero_point)

def get_inputs():
    x = torch.randn(SHAPE, dtype=torch.float32)
    scale = torch.tensor([0.05], dtype=torch.float32)
    zero_point = torch.tensor([10], dtype=torch.int32)
    
    return [x.contiguous(), scale, zero_point]

def get_init_inputs():
    return []