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

python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py

@@ -0,0 +1,267 @@
+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 == "BatchNormalization":
+                pytest.skip(f"{node.op_type} alters the node structure by design")
+            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}")

python/tests/onnx_quantizer_tests/test_base_layer.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import onnx
+import pytest
+from onnx import helper, numpy_helper
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from onnx import GraphProto, ModelProto, NodeProto, TensorProto
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    HandlerImplementationError,
+)
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    ScaleConfig,
+)
+
+
+class DummyQuantizer(BaseOpQuantizer):
+    def __init__(self: DummyQuantizer) -> None:
+        self.new_initializers = []
+
+
+@pytest.fixture
+def dummy_tensor() -> TensorProto:
+    return numpy_helper.from_array(np.array([[1.0, 2.0], [3.0, 4.0]]), name="W")
+
+
+@pytest.fixture
+def dummy_bias() -> TensorProto:
+    return numpy_helper.from_array(np.array([1.0, 2.0]), name="B")
+
+
+@pytest.fixture
+def dummy_node() -> NodeProto:
+    return helper.make_node(
+        "DummyOp",
+        inputs=["X", "W", "B"],
+        outputs=["Y"],
+        name="DummyOp",
+    )
+
+
+@pytest.fixture
+def dummy_graph() -> GraphProto:
+    return helper.make_graph([], "dummy_graph", inputs=[], outputs=[])
+
+
+@pytest.fixture
+def initializer_map(
+    dummy_tensor: TensorProto,
+    dummy_bias: TensorProto,
+) -> dict[str, TensorProto]:
+    return {"W": dummy_tensor, "B": dummy_bias}
+
+
+@pytest.fixture
+def minimal_model() -> ModelProto:
+    graph = onnx.helper.make_graph(
+        nodes=[],  # No nodes
+        name="EmptyGraph",
+        inputs=[],
+        outputs=[],
+        initializer=[],
+    )
+    return onnx.helper.make_model(graph)
+
+
+@pytest.fixture
+def unsupported_model() -> ModelProto:
+    node = onnx.helper.make_node("UnsupportedOp", ["X"], ["Y"])
+    graph = onnx.helper.make_graph(
+        nodes=[node],
+        name="UnsupportedGraph",
+        inputs=[],
+        outputs=[],
+        initializer=[],
+    )
+    return onnx.helper.make_model(graph)
+
+
+@pytest.mark.unit
+def test_quantize_raises_not_implemented() -> None:
+    quantizer = BaseOpQuantizer()
+    with pytest.raises(
+        HandlerImplementationError,
+    ) as excinfo:
+        quantizer.quantize(
+            node=None,
+            graph=None,
+            scale_config=ScaleConfig(exponent=1, base=1, rescale=False),
+            initializer_map={},
+        )
+    assert "quantize() not implemented in subclass." in str(excinfo.value)
+
+
+@pytest.mark.unit
+def test_check_supported_returns_none(dummy_node: NodeProto) -> None:
+    quantizer = DummyQuantizer()
+    with pytest.raises(HandlerImplementationError) as excinfo:
+        quantizer.check_supported(dummy_node, {})
+
+    assert (
+        "Handler implementation error for operator 'DummyQuantizer':"
+        " check_supported() not implemented in subclass." in str(excinfo.value)
+    )
+
+
+@pytest.mark.unit
+def test_rescale_layer_modifies_node_output(
+    dummy_node: NodeProto,
+    dummy_graph: GraphProto,
+) -> None:
+    quantizer = DummyQuantizer()
+    result_nodes = quantizer.rescale_layer(
+        dummy_node,
+        scale_base=10,
+        scale_exponent=2,
+        graph=dummy_graph,
+    )
+    total_scale = 100.0
+    count_nodes = 2
+
+    assert len(result_nodes) == count_nodes
+    assert dummy_node.output[0] == "Y_raw"
+    assert result_nodes[1].op_type == "Div"
+    assert result_nodes[1].output[0] == "Y"
+
+    # Check if scale tensor added
+    assert len(quantizer.new_initializers) == 1
+    scale_tensor = quantizer.new_initializers[0]
+    assert scale_tensor.name.endswith("_scale")
+    assert scale_tensor.data_type == onnx.TensorProto.INT64
+    assert onnx.numpy_helper.to_array(scale_tensor)[0] == total_scale
+
+    # Validate that result_nodes are valid ONNX nodes
+    for node in result_nodes:
+        assert isinstance(node, onnx.NodeProto)
+        assert node.name
+        assert node.op_type
+        assert node.input
+        assert node.output
+
+    # Check Div node inputs: should divide Y_raw by scale
+    div_node = result_nodes[1]
+    assert len(div_node.input) == count_nodes
+    assert div_node.input[0] == "Y_raw"
+    assert div_node.input[1] == scale_tensor.name
+
+
+@pytest.mark.unit
+def test_add_nodes_w_and_b_creates_mul_and_cast(
+    dummy_node: NodeProto,
+    dummy_graph: GraphProto,
+    initializer_map: dict[str, Any],
+) -> None:
+    _ = dummy_graph
+    quantizer = DummyQuantizer()
+    exp = 2
+    base = 10
+    nodes, new_inputs = quantizer.add_nodes_w_and_b(
+        dummy_node,
+        scale_exponent=exp,
+        scale_base=base,
+        initializer_map=initializer_map,
+    )
+    four = 4
+    two = 2
+
+    assert len(nodes) == four  # Mul + Cast for W, Mul + Cast for B
+    assert nodes[0].op_type == "Mul"
+    assert nodes[1].op_type == "Cast"
+    assert nodes[2].op_type == "Mul"
+    assert nodes[3].op_type == "Cast"
+    assert new_inputs == ["X", "W_scaled_cast", "B_scaled_cast"]
+    assert len(quantizer.new_initializers) == two
+
+    weight_scaled = base**exp
+    bias_scaled = base ** (exp * 2)
+
+    # Enhanced assertions: check node inputs/outputs and tensor details
+    # Mul for W: input W and W_scale, output W_scaled
+    assert nodes[0].input == ["W", "W_scale"]
+    assert nodes[0].output == ["W_scaled"]
+    # Cast for W: input W_scaled, output W_scaled_cast
+    assert nodes[1].input == ["W_scaled"]
+    assert nodes[1].output == ["W_scaled_cast"]
+    # Similarly for B
+    assert nodes[2].input == ["B", "B_scale"]
+    assert nodes[2].output == ["B_scaled"]
+    assert nodes[3].input == ["B_scaled"]
+    assert nodes[3].output == ["B_scaled_cast"]
+
+    # Check scale tensors
+    w_scale = quantizer.new_initializers[0]
+    b_scale = quantizer.new_initializers[1]
+    assert w_scale.name == "W_scale"
+    assert b_scale.name == "B_scale"
+    assert onnx.numpy_helper.to_array(w_scale)[0] == weight_scaled  # 10**2
+    assert onnx.numpy_helper.to_array(b_scale)[0] == bias_scaled
+
+
+@pytest.mark.unit
+def test_insert_scale_node_creates_mul_and_cast(
+    dummy_tensor: TensorProto,
+    dummy_graph: GraphProto,
+) -> None:
+    _ = dummy_graph
+    quantizer = DummyQuantizer()
+    output_name, mul_node, cast_node = quantizer.insert_scale_node(
+        dummy_tensor,
+        scale_base=10,
+        scale_exponent=1,
+    )
+
+    assert mul_node.op_type == "Mul"
+    assert cast_node.op_type == "Cast"
+    assert "_scaled" in mul_node.output[0]
+    assert output_name.endswith("_cast")
+    assert len(quantizer.new_initializers) == 1
+    assert quantizer.new_initializers[0].name.endswith("_scale")
+    ten = 10.0
+    assert onnx.numpy_helper.to_array(quantizer.new_initializers[0])[0] == ten