import torch
import torch.nn as nn

BATCH_SIZE = 64
CHANNELS = 64
HEIGHT = 64
WIDTH = 64
SHAPE = (BATCH_SIZE, CHANNELS, HEIGHT, WIDTH)

AXIS = 1
Q_MIN = -128
Q_MAX = 127

class Model(nn.Module):
    def __init__(self, axis, q_min, q_max):
        super(Model, self).__init__()
        self.axis = axis
        self.q_min = q_min
        self.q_max = q_max

    def forward(self, x, scale, zero_point):
        return torch.fake_quantize_per_channel_affine(
            x, scale, zero_point, self.axis, self.q_min, self.q_max
        )

def get_inputs():
    x = torch.randn(SHAPE, dtype=torch.float32)
    # Per-channel parameters
    scale = torch.rand((CHANNELS,), dtype=torch.float32)
    zero_point = torch.randint(0, 10, (CHANNELS,), dtype=torch.int32)
    
    return [x.contiguous(), scale.contiguous(), zero_point.contiguous()]

def get_init_inputs():
    return [AXIS, Q_MIN, Q_MAX]