def test_dpu_export_onnx_quant_conv_bias():
    FEATURES = 7
    IN_SIZE = (1, IN_CH, FEATURES, FEATURES)
    KERNEL_SIZE = 3

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = QuantConv2d(
                out_channels=OUT_CH,
                in_channels=IN_CH,
                kernel_size=KERNEL_SIZE,
                bias=True,
                weight_quant=Int8WeightPerTensorFixedPoint,
                bias_quant=Int8BiasPerTensorFixedPointInternalScaling,
                input_quant=Int8ActPerTensorFixedPoint,
                output_quant=Int8ActPerTensorFixedPoint,
                return_quant_tensor=False)
            self.conv1.weight.data.uniform_(-0.01, 0.01)

        def forward(self, x):
            return self.conv1(x)

    inp = gen_linspaced_data(reduce(mul, IN_SIZE), -1, 1).reshape(IN_SIZE)
    model = Model()
    model(torch.from_numpy(inp))  # accumulate scale factors
    model.eval()
    PyXIRManager.export(model, input_shape=IN_SIZE)
def test_standard_onnx_quant_linear_bias_export():
    IN_SIZE = (IN_CH, IN_CH)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = QuantLinear(
                in_features=IN_CH,
                out_features=OUT_CH,
                bias=True,
                weight_quant=Int8WeightPerTensorFixedPoint,
                bias_quant=Int8BiasPerTensorFixedPointInternalScaling,
                input_quant=Int8ActPerTensorFixedPoint,
                output_quant=Int8ActPerTensorFixedPoint,
                return_quant_tensor=False)
            self.linear.weight.data.uniform_(-0.01, 0.01)

        def forward(self, x):
            return self.linear(x)

    inp = gen_linspaced_data(reduce(mul, IN_SIZE)).reshape(IN_SIZE)
    model = Model()
    model(torch.from_numpy(inp))  # accumulate scale factors
    model.eval()
    PyXIRManager.export(model, input_shape=IN_SIZE)