JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl → 1.1.0__py3-none-macosx_11_0_arm64.whl

python/tests/onnx_quantizer_tests/layers_tests/test_integration.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+from onnxruntime import InferenceSession, SessionOptions
+from onnxruntime_extensions import get_library_path
+
+from python.core.model_processing.converters.onnx_converter import ONNXConverter
+
+if TYPE_CHECKING:
+    from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (
+        ONNXOpQuantizer,
+    )
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+from python.tests.onnx_quantizer_tests.layers.base import (
+    LayerTestConfig,
+    LayerTestSpec,
+    SpecType,
+)
+from python.tests.onnx_quantizer_tests.layers.factory import TestLayerFactory
+from python.tests.onnx_quantizer_tests.layers_tests.base_test import (
+    BaseQuantizerTest,
+)
+
+
+class TestIntegration(BaseQuantizerTest):
+    """Integration tests for ONNX quantizer"""
+
+    __test__ = True
+
+    @pytest.mark.integration
+    @pytest.mark.parametrize(
+        "layer_combination",
+        [["Conv", "Relu"], ["Gemm", "Relu"], ["Conv", "MaxPool", "Flatten", "Gemm"]],
+    )
+    def test_check_then_quantize_workflow(
+        self: TestIntegration,
+        quantizer: ONNXOpQuantizer,
+        layer_configs: dict[str, LayerTestConfig],
+        layer_combination: list[str],
+    ) -> None:
+        """Test the typical workflow: check model then quantize layers"""
+        mock_graph = Mock()
+        scale_exponent, scale_base = 2, 10
+        rescale = True
+
+        # Step 1: Create and validate model
+        model = self.create_model_with_layers(layer_combination, layer_configs)
+        quantizer.check_model(model)  # Should not raise
+
+        # Step 2: Quantize each layer
+        initializer_map = quantizer.get_initializer_map(model)
+
+        for node in model.graph.node:
+            result = quantizer.quantize(
+                node=node,
+                rescale=rescale,
+                graph=mock_graph,
+                scale_exponent=scale_exponent,
+                scale_base=scale_base,
+                initializer_map=initializer_map,
+            )
+            assert result is not None, (
+                f"Quantization failed for {node.op_type}"
+                f" in combination {layer_combination}"
+            )
+
+    def skip_by_layer_name(
+        self,
+        layer_name: str,
+        test_spec: LayerTestSpec,
+        skip_layer: str,
+    ) -> None:
+        # Skip Constant nodes as they don't depend on scaled inputs
+        if layer_name == skip_layer:
+            pytest.skip(
+                f"Skipping accuracy test for {layer_name}."
+                f"{test_spec.name} as constants are scaled differently",
+            )
+
+    @pytest.mark.integration
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.VALID),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_end_to_end_quantization_accuracy(
+        self: TestIntegration,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        """Test end-to-end quantization accuracy for each valid test case.
+
+        Builds a model from the layer config, runs inference on the original model,
+        quantizes the model, runs inference on the quantized model, and ensures
+        the outputs are close.
+        """
+        cosine_similarity = 0.995
+        rng = np.random.default_rng(TEST_RNG_SEED + 1)
+
+        layer_name, config, test_spec = test_case_data
+        self.skip_by_layer_name(layer_name, test_spec, skip_layer="Constant")
+
+        # Skip if validation failed or test is skipped
+        self._check_validation_dependency(test_case_data)
+        if test_spec.skip_reason:
+            pytest.skip(f"{layer_name}_{test_spec.name}: {test_spec.skip_reason}")
+
+        # Create original model
+        original_model = config.create_test_model(test_spec)
+        opts = SessionOptions()
+        opts.register_custom_ops_library(get_library_path())
+        original_session = InferenceSession(
+            original_model.SerializeToString(),
+            opts,
+            providers=["CPUExecutionProvider"],
+        )
+
+        input_shapes = {i.name: tuple(i.shape) for i in original_session.get_inputs()}
+
+        # Skip if no inputs (e.g., Constant nodes)
+        if not input_shapes:
+            pytest.skip(
+                f"No inputs for {layer_name}.{test_spec.name}, skipping accuracy test",
+            )
+
+        # Create dummy inputs for all graph inputs
+
+        dummy_inputs = {}
+        for name, shape in input_shapes.items():
+            dummy_inputs[name] = rng.normal(0, 1, shape).astype(np.float32)
+
+        # Run inference on original model
+        output_name = original_session.get_outputs()[0].name
+        original_output = original_session.run([output_name], dummy_inputs)[0]
+
+        # Quantize the model
+
+        converter = ONNXConverter()
+        scale_base, scale_exponent = (
+            2,
+            10,
+        )  # Smaller scale to reduce quantization errors
+        quantized_model = converter.quantize_model(
+            original_model,
+            scale_base=scale_base,
+            scale_exponent=scale_exponent,
+            rescale_config=None,  # Use default rescale
+        )
+
+        # Run inference on quantized model
+        quantized_session = InferenceSession(
+            quantized_model.SerializeToString(),
+            opts,
+            providers=["CPUExecutionProvider"],
+        )
+        quantized_input_names = [inp.name for inp in quantized_session.get_inputs()]
+        quantized_output_name = quantized_session.get_outputs()[0].name
+
+        # For quantized model, scale the inputs
+        scaled_inputs = {}
+        for name in quantized_input_names:
+            if name in dummy_inputs:
+                scaled_inputs[name] = (dummy_inputs[name]).astype(np.float64)
+            else:
+                # If quantized model has different inputs, skip or handle
+                pytest.skip(
+                    f"Quantized model input mismatch for {layer_name}.{test_spec.name}",
+                )
+
+        quantized_output = quantized_session.run(
+            [quantized_output_name],
+            scaled_inputs,
+        )[0]
+        quantized_output = quantized_output / (scale_base ** (scale_exponent))
+
+        ratio = np.mean(quantized_output / (original_output + 1e-12))
+        print(f"Mean output ratio (quantized/original): {ratio:.4f}")
+
+        # Compare outputs (quantized output should be close to original if rescale=True)
+        # Allow some tolerance due to quantization
+        np.testing.assert_allclose(
+            original_output,
+            quantized_output,
+            rtol=0.05,  # Relative tolerance
+            atol=0.05,  # Absolute tolerance
+            err_msg=f"Quantization accuracy failed for {layer_name}.{test_spec.name}",
+        )
+
+        cos_sim = np.dot(original_output.flatten(), quantized_output.flatten()) / (
+            np.linalg.norm(original_output.flatten())
+            * np.linalg.norm(quantized_output.flatten())
+            + 1e-12
+        )
+        print(f"Cosine similarity: {cos_sim:.6f}")
+        assert cos_sim > cosine_similarity, f"Low cosine similarity ({cos_sim:.6f})"

python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from unittest.mock import Mock
+
+import onnx
+import pytest
+from onnx import NodeProto
+
+from python.core.model_processing.onnx_quantizer.exceptions import UnsupportedOpError
+
+if TYPE_CHECKING:
+    from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (
+        ONNXOpQuantizer,
+    )
+from python.tests.onnx_quantizer_tests.layers.base import (
+    LayerTestConfig,
+    LayerTestSpec,
+    SpecType,
+)
+from python.tests.onnx_quantizer_tests.layers.factory import TestLayerFactory
+from python.tests.onnx_quantizer_tests.layers_tests.base_test import (
+    BaseQuantizerTest,
+)
+
+
+class TestQuantize(BaseQuantizerTest):
+    """Tests for quantization functionality"""
+
+    __test__ = True
+
+    def setup_quantize_test(
+        self,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+        quantizer: ONNXOpQuantizer,
+        scale_exponent: int = 2,
+        scale_base: int = 10,
+        *,
+        rescale: bool = True,
+    ) -> tuple[onnx.NodeProto, tuple[str, LayerTestConfig, LayerTestSpec, NodeProto]]:
+        """Common setup for quantization tests"""
+        layer_name, config, test_spec = test_case_data
+
+        self._check_validation_dependency(test_case_data)
+
+        if test_spec.skip_reason:
+            pytest.skip(f"{layer_name}_{test_spec.name}: {test_spec.skip_reason}")
+
+        model = config.create_test_model(test_spec)
+        node = model.graph.node[0]
+        initializer_map = {init.name: init for init in model.graph.initializer}
+
+        mock_graph = Mock()
+        if node.op_type == "Constant":
+            mock_data_node = Mock()
+            mock_data_node.input = [node.output[0]]
+            mock_graph.node = [mock_data_node]
+
+        result = quantizer.quantize(
+            node=node,
+            rescale=rescale,
+            graph=mock_graph,
+            scale_exponent=scale_exponent,
+            scale_base=scale_base,
+            initializer_map=initializer_map,
+        )
+
+        return result, (layer_name, config, test_spec, node)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.VALID),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_quantize_individual_valid_cases(
+        self: TestQuantize,
+        quantizer: ONNXOpQuantizer,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        """Test quantization for each individual valid test case"""
+
+        scale_exponent, scale_base = 2, 10
+        rescale = True
+
+        result, (layer_name, _config, test_spec, _node) = self.setup_quantize_test(
+            test_case_data,
+            quantizer,
+            scale_exponent,
+            scale_base,
+            rescale=rescale,
+        )
+
+        # Test that the output of the quantizer quantize is in fact a node
+        if isinstance(result, list):
+            assert (
+                len(result) > 0
+            ), f"Quantize returned empty list for {layer_name}.{test_spec.name}"
+            for node_result in result:
+                assert isinstance(
+                    node_result,
+                    onnx.NodeProto,
+                ), f"Invalid node type returned for {layer_name}.{test_spec.name}"
+        else:
+            assert isinstance(
+                result,
+                onnx.NodeProto,
+            ), f"Quantize returned none node for {layer_name}.{test_spec.name}"
+
+            assert (
+                result is not None
+            ), f"Quantize returned None for {layer_name}.{test_spec.name}"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.VALID),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_quantize_preserves_node_names(
+        self: TestQuantize,
+        quantizer: ONNXOpQuantizer,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        """Test quantization for each individual valid test case"""
+
+        scale_exponent, scale_base = 2, 10
+        rescale = True
+        result, (_layer_name, config, _test_spec, node) = self.setup_quantize_test(
+            test_case_data,
+            quantizer,
+            scale_exponent,
+            scale_base,
+            rescale=rescale,
+        )
+        is_node_present = False
+
+        def check_node_and_analyze_parameters(
+            node: NodeProto,
+            result_node: NodeProto,
+        ) -> bool:
+            if node.op_type in result_node.op_type:
+                # Assert there are no less attributes in the new node
+                assert len(node.attribute) <= len(result_node.attribute)
+                # Ensure that each original node's attributes
+                # are contained in the new nodes
+                for att in node.attribute:
+                    assert att.name in [a.name for a in result_node.attribute]
+                return True
+            return False
+
+        # Check that result nodes have meaningful names and the relevant node is present
+        # And ensure that the new node has the same parameters as the old node
+        if isinstance(result, list):
+            for result_node in result:
+                assert (
+                    result_node.name
+                ), f"Quantized node missing name for {config.op_type}"
+                assert (
+                    result_node.op_type
+                ), f"Quantized node missing op_type for {config.op_type}"
+
+                is_node_present = is_node_present or check_node_and_analyze_parameters(
+                    node,
+                    result_node,
+                )
+        else:
+            assert result.name, f"Quantized node missing name for {config.op_type}"
+            is_node_present = is_node_present or check_node_and_analyze_parameters(
+                node,
+                result,
+            )
+
+        # Assert that the node is in fact present
+        assert (
+            is_node_present
+        ), "Cannot find quantized node relating to prequantized node"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("scale_params", [(2, 10), (0, 5)])
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.VALID),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_quantize_with_different_scales(
+        self: TestQuantize,
+        quantizer: ONNXOpQuantizer,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+        scale_params: tuple[int, int],
+    ) -> None:
+        """Test quantization for each individual valid test case"""
+
+        # Test for both scale parameters
+        scale_exponent, scale_base = scale_params
+        rescale = True
+        result, (_layer_name, _config, _test_spec, _node) = self.setup_quantize_test(
+            test_case_data,
+            quantizer,
+            scale_exponent,
+            scale_base,
+            rescale=rescale,
+        )
+
+        # Should return valid result regardless of scale values
+        assert (
+            result is not None
+        ), f"Quantize returned None for scale={scale_exponent}, scale_base={scale_base}"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("rescale", [True, False])
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.VALID),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_quantize_with_different_rescales(
+        self: TestQuantize,
+        quantizer: ONNXOpQuantizer,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+        *,
+        rescale: bool,
+    ) -> None:
+        """Test quantization for each individual valid test case"""
+
+        scale_exponent, scale_base = 2, 10
+
+        # Test that quantizing works with both rescaling values
+        result, (_layer_name, _config, _test_spec, _node) = self.setup_quantize_test(
+            test_case_data,
+            quantizer,
+            scale_exponent,
+            scale_base,
+            rescale=rescale,
+        )
+        assert result is not None, f"Quantize failed with rescale={rescale}"
+
+    @pytest.mark.unit
+    def test_quantize_unsupported_layer_returns_original(
+        self: TestQuantize,
+        quantizer: ONNXOpQuantizer,
+    ) -> None:
+        """Test that unsupported layers return Error in quantization process"""
+        from onnx import helper  # noqa: PLC0415
+
+        mock_graph = Mock()
+        scale_exponent, scale_base = 2, 10
+        rescale = True
+
+        unsupported_node = helper.make_node(
+            "UnsupportedOp",
+            inputs=["input"],
+            outputs=["output"],
+            name="unsupported",
+        )
+        with pytest.raises(UnsupportedOpError) as excinfo:
+            _ = quantizer.quantize(
+                node=unsupported_node,
+                rescale=rescale,
+                graph=mock_graph,
+                scale_exponent=scale_exponent,
+                scale_base=scale_base,
+                initializer_map={},
+            )
+        assert "Unsupported op type: 'UnsupportedOp'" in str(excinfo.value)

python/tests/onnx_quantizer_tests/layers_tests/test_scalability.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (
+    ONNXOpQuantizer,
+)
+from python.tests.onnx_quantizer_tests.layers.base import LayerTestConfig, SpecType
+from python.tests.onnx_quantizer_tests.layers.factory import TestLayerFactory
+
+
+class TestScalability:
+    """Tests (meta) to verify the framework scales with new layers"""
+
+    @pytest.mark.unit
+    def test_adding_new_layer_config(self: TestScalability) -> None:
+        """Test that adding new layer configs is straightforward"""
+        two = 2
+        # Simulate adding a new layer type
+        new_layer_config = LayerTestConfig(
+            op_type="NewCustomOp",
+            valid_inputs=["input", "custom_param"],
+            valid_attributes={"custom_attr": 42},
+            required_initializers={"custom_param": np.array([1, 2, 3])},
+        )
+
+        # Verify config can create nodes and initializers
+        node = new_layer_config.create_node()
+        assert node.op_type == "NewCustomOp"
+        assert len(node.input) == two
+
+        initializers = new_layer_config.create_initializers()
+        assert "custom_param" in initializers
+
+    @pytest.mark.unit
+    def test_layer_config_extensibility(self: TestScalability) -> None:
+        """Test that layer configs consists of all registered handlers"""
+        configs = TestLayerFactory.get_layer_configs()
+
+        # Verify all expected layers are present
+        unsupported = ONNXOpQuantizer().handlers.keys() - set(configs.keys())
+        assert unsupported == set(), (
+            f"The following layers are not being configured for testing: {unsupported}."
+            " Please add configuration in tests/onnx_quantizer_tests/layers/"
+        )
+
+        # Verify each config has required components
+        for layer_type, config in configs.items():
+            err_msg = (
+                f"Quantization test config is not supported yet for {layer_type}"
+                " and must be implemented"
+            )
+            assert config.op_type == layer_type, err_msg
+            assert isinstance(
+                config.valid_inputs,
+                list,
+            ), err_msg
+            assert isinstance(
+                config.valid_attributes,
+                dict,
+            ), err_msg
+            assert isinstance(
+                config.required_initializers,
+                dict,
+            ), err_msg
+
+    @pytest.mark.unit
+    def test_every_layer_has_basic_and_e2e(self: TestScalability) -> None:
+        """Each registered layer must have at least one basic/valid test
+        and one e2e test."""
+        missing_basic = []
+        missing_e2e = []
+
+        # iterate over registered layers
+        for layer_name in TestLayerFactory.get_available_layers():
+            cases = TestLayerFactory.get_test_cases_by_layer(layer_name)
+            specs = [spec for _, _config, spec in cases]
+
+            # Consider a test "basic" if:
+            #  - it has tag 'basic' or 'valid', OR
+            #  - its spec_type is SpecType.VALID (if you use SpecType)
+            has_basic = any(
+                (
+                    "basic" in getattr(s, "tags", set())
+                    or "valid" in getattr(s, "tags", set())
+                    or getattr(s, "spec_type", None) == SpecType.VALID
+                )
+                for s in specs
+            )
+
+            # Consider a test "e2e" if:
+            #  - it has tag 'e2e', OR
+            #  - its spec_type is SpecType.E2E (if you use that enum)
+            has_e2e = any(
+                (
+                    "e2e" in getattr(s, "tags", set())
+                    or getattr(s, "spec_type", None) == SpecType.E2E
+                )
+                for s in specs
+            )
+
+            if not has_basic:
+                missing_basic.append(layer_name)
+            if not has_e2e:
+                missing_e2e.append(layer_name)
+
+        assert not missing_basic, f"Layers missing a basic/valid test: {missing_basic}"
+        assert not missing_e2e, f"Layers missing an e2e test: {missing_e2e}"

python/tests/onnx_quantizer_tests/layers_tests/test_validation.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import onnx
+import pytest
+
+if TYPE_CHECKING:
+    from python.tests.onnx_quantizer_tests.layers.base import (
+        LayerTestConfig,
+        LayerTestSpec,
+    )
+from python.tests.onnx_quantizer_tests.layers.factory import TestLayerFactory
+from python.tests.onnx_quantizer_tests.layers_tests.base_test import (
+    BaseQuantizerTest,
+)
+
+
+class TestValidation(BaseQuantizerTest):
+    """Ensure that layer factory models produce valid ONNX graphs."""
+
+    __test__ = True
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_all_test_cases(),
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_factory_models_pass_onnx_validation(
+        self: TestValidation,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        layer_name, config, test_spec = test_case_data
+        test_case_id = f"{layer_name}_{test_spec.name}"
+
+        if test_spec.skip_reason:
+            pytest.skip(f"{test_case_id}: {test_spec.skip_reason}")
+
+        model = config.create_test_model(test_spec)
+        try:
+            onnx.checker.check_model(model)
+        except onnx.checker.ValidationError as e:
+            self._validation_failed_cases.add(test_case_id)
+            pytest.fail(f"Invalid ONNX model: {e}")