diff --git a/test/quantization/test_algorithm_utility.py b/test/quantization/test_algorithm_utility.py index 586d851d3..28a8f2a7a 100644 --- a/test/quantization/test_algorithm_utility.py +++ b/test/quantization/test_algorithm_utility.py @@ -1,16 +1,17 @@ """Tests for algorithm utility components.""" import os -import onnx import shutil import unittest + import numpy as np +import onnx import onnxruntime -import optimum.exporters.onnx import onnxruntime.tools.symbolic_shape_infer as symbolic_shape_infer +import optimum.exporters.onnx -from onnx_neural_compressor.algorithms import utility as quant_utils from onnx_neural_compressor import onnx_model +from onnx_neural_compressor.algorithms import utility as quant_utils def find_onnx_file(folder_path): @@ -21,6 +22,7 @@ def find_onnx_file(folder_path): return os.path.join(root, file) return None + class TestUtilityFunctions(unittest.TestCase): @classmethod def setUpClass(self): @@ -77,7 +79,9 @@ def test_quantize_data(self): scale, zero_point = quant_utils._calculate_scale_zp(np.array([0]), np.array([5]), quantize_range, qType, scheme) self.assertEqual(zero_point.dtype, np.int8) - scale, zero_point = quant_utils._calculate_scale_zp(np.array([0]), np.array([127]), quantize_range, qType, scheme) + scale, zero_point = quant_utils._calculate_scale_zp( + np.array([0]), np.array([127]), quantize_range, qType, scheme + ) self.assertEqual(zero_point.dtype, np.int8) # asym uint8 @@ -91,7 +95,9 @@ def test_quantize_data(self): scale, zero_point = quant_utils._calculate_scale_zp(np.array([0]), np.array([5]), quantize_range, qType, scheme) self.assertEqual(zero_point.dtype, np.uint8) - scale, zero_point = quant_utils._calculate_scale_zp(np.array([0]), np.array([255]), quantize_range, qType, scheme) + scale, zero_point = quant_utils._calculate_scale_zp( + np.array([0]), np.array([255]), quantize_range, qType, scheme + ) self.assertEqual(zero_point.dtype, np.uint8) # unexpected combination