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/test_exceptions.py

@@ -0,0 +1,99 @@
+import pytest
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    REPORTING_URL,
+    InvalidParamError,
+    QuantizationError,
+    UnsupportedOpError,
+)
+
+
+@pytest.mark.unit
+def test_quantization_error_message() -> None:
+    custom_msg = "Something went wrong."
+    with pytest.raises(QuantizationError) as exc_info:
+        raise QuantizationError(custom_msg)
+    assert "This model is not supported by JSTprove." in str(exc_info.value)
+    assert custom_msg in str(exc_info.value)
+
+    assert REPORTING_URL in str(exc_info.value)
+
+    assert "Submit model support requests via the JSTprove channel:" in str(
+        exc_info.value,
+    )
+
+
+@pytest.mark.unit
+def test_invalid_param_error_basic() -> None:
+    with pytest.raises(InvalidParamError) as exc_info:
+        raise InvalidParamError(
+            node_name="Conv_1",
+            op_type="Conv",
+            message="Missing 'strides' attribute.",
+        )
+    err_msg = str(exc_info.value)
+    assert "Invalid parameters in node 'Conv_1'" in err_msg
+    assert "(op_type='Conv')" in err_msg
+    assert "Missing 'strides' attribute." in err_msg
+    assert "[Attribute:" not in err_msg
+    assert "[Expected:" not in err_msg
+
+    msg = ""
+
+    with pytest.raises(QuantizationError) as exc_info_quantization:
+        raise QuantizationError(msg)
+    # Assert contains generic error message from quantization error
+    assert str(exc_info_quantization.value) in err_msg
+
+
+@pytest.mark.unit
+def test_invalid_param_error_with_attr_and_expected() -> None:
+    with pytest.raises(InvalidParamError) as exc_info:
+        raise InvalidParamError(
+            node_name="MaxPool_3",
+            op_type="MaxPool",
+            message="Kernel shape is invalid.",
+            attr_key="kernel_shape",
+            expected="a list of 2 positive integers",
+        )
+    err_msg = str(exc_info.value)
+    assert "Invalid parameters in node 'MaxPool_3'" in err_msg
+    assert "[Attribute: kernel_shape]" in err_msg
+    assert "[Expected: a list of 2 positive integers]" in err_msg
+    msg = ""
+
+    with pytest.raises(QuantizationError) as exc_info_quantization:
+        raise QuantizationError(msg)
+    # Assert contains generic error message from quantization error
+    assert str(exc_info_quantization.value) in err_msg
+
+
+@pytest.mark.unit
+def test_unsupported_op_error_with_node() -> None:
+    with pytest.raises(UnsupportedOpError) as exc_info:
+        raise UnsupportedOpError(op_type="Resize", node_name="Resize_42")
+    err_msg = str(exc_info.value)
+    assert "Unsupported op type: 'Resize'" in err_msg
+    assert "in node 'Resize_42'" in err_msg
+    assert "documentation for supported layers" in err_msg
+    msg = ""
+
+    with pytest.raises(QuantizationError) as exc_info_quantization:
+        raise QuantizationError(msg)
+    # Assert contains generic error message from quantization error
+    assert str(exc_info_quantization.value) in err_msg
+
+
+@pytest.mark.unit
+def test_unsupported_op_error_without_node() -> None:
+    with pytest.raises(UnsupportedOpError) as exc_info:
+        raise UnsupportedOpError(op_type="Upsample")
+    err_msg = str(exc_info.value)
+    assert "Unsupported op type: 'Upsample'" in err_msg
+    assert "in node" not in err_msg
+    msg = ""
+
+    with pytest.raises(QuantizationError) as exc_info_quantization:
+        raise QuantizationError(msg)
+    # Assert contains generic error message from quantization error
+    assert str(exc_info_quantization.value) in err_msg

python/tests/onnx_quantizer_tests/test_onnx_op_quantizer.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import onnx
+import pytest
+from onnx import GraphProto, ModelProto, NodeProto, TensorProto, helper
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    MissingHandlerError,
+    QuantizationError,
+    UnsupportedOpError,
+)
+
+if TYPE_CHECKING:
+    from python.core.model_processing.onnx_quantizer.layers.base import ScaleConfig
+
+# Optional: mock layers if needed
+from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (
+    ONNXOpQuantizer,
+)
+
+
+# Mocks
+class MockHandler:
+    def __init__(self: MockHandler) -> None:
+        self.called_quantize = False
+        self.called_supported = False
+
+    def quantize(
+        self: MockHandler,
+        node: NodeProto,
+        graph: GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, TensorProto],
+    ) -> list[NodeProto]:
+        _ = graph, scale_config, initializer_map
+        self.called_quantize = True
+        return [node]  # Return the original node as a list for simplicity
+
+    def check_supported(
+        self: MockHandler,
+        node: NodeProto,
+        initializer_map: dict[str, TensorProto],
+    ) -> None:
+        _ = initializer_map
+        self.called_supported = True
+        if node.name == "bad_node":
+            msg = "Invalid node parameters"
+            raise ValueError(msg)
+
+
+# Fixtures
+@pytest.fixture
+def quantizer() -> ONNXOpQuantizer:
+    return ONNXOpQuantizer()
+
+
+@pytest.fixture
+def dummy_node() -> NodeProto:
+    return helper.make_node("FakeOp", inputs=["x"], outputs=["y"])
+
+
+@pytest.fixture
+def valid_node() -> NodeProto:
+    return helper.make_node("Dummy", inputs=["x"], outputs=["y"], name="good_node")
+
+
+@pytest.fixture
+def invalid_node() -> NodeProto:
+    return helper.make_node("Dummy", inputs=["x"], outputs=["y"], name="bad_node")
+
+
+@pytest.fixture
+def dummy_model(valid_node: NodeProto, invalid_node: NodeProto) -> ModelProto:
+    graph = helper.make_graph(
+        [valid_node, invalid_node],
+        "test_graph",
+        inputs=[],
+        outputs=[],
+        initializer=[helper.make_tensor("x", TensorProto.FLOAT, [1], [0.5])],
+    )
+    return helper.make_model(graph)
+
+
+# Tests
+
+
+@pytest.mark.unit
+def test_check_model_raises_on_unsupported_op() -> None:
+    quantizer = ONNXOpQuantizer()
+
+    unsupported_node = helper.make_node("UnsupportedOp", ["x"], ["y"])
+    graph = helper.make_graph([unsupported_node], "test_graph", [], [])
+    model = helper.make_model(graph)
+
+    with pytest.raises(UnsupportedOpError):
+        quantizer.check_model(model)
+
+
+@pytest.mark.unit
+def test_check_layer_invokes_check_supported() -> None:
+    quantizer = ONNXOpQuantizer()
+    handler = MockHandler()
+    quantizer.register("FakeOp", handler)
+
+    node = helper.make_node("FakeOp", ["x"], ["y"])
+    initializer_map = {}
+
+    quantizer.check_layer(node, initializer_map)
+    # Check that check_supported is called
+    assert handler.called_supported
+
+
+@pytest.mark.unit
+def test_get_initializer_map_returns_correct_dict() -> None:
+    quantizer = ONNXOpQuantizer()
+
+    tensor = helper.make_tensor(
+        name="W",
+        data_type=TensorProto.FLOAT,
+        dims=[1],
+        vals=[1.0],
+    )
+    graph = helper.make_graph([], "test_graph", [], [], [tensor])
+    model = helper.make_model(graph)
+
+    init_map = quantizer.get_initializer_map(model)
+    # Test initializer in map
+    assert "W" in init_map
+    # Test initializer map lines up
+    assert init_map["W"] == tensor
+    # Enhanced: check tensor properties
+    assert init_map["W"].data_type == TensorProto.FLOAT
+    assert init_map["W"].dims == [1]
+    assert onnx.numpy_helper.to_array(init_map["W"])[0] == 1.0
+
+
+@pytest.mark.unit
+def test_quantize_with_unregistered_op_warns(dummy_node: NodeProto) -> None:
+    quantizer = ONNXOpQuantizer()
+    graph = helper.make_graph([], "g", [], [])
+    with pytest.raises(UnsupportedOpError) as excinfo:
+        _ = quantizer.quantize(dummy_node, graph, 1, 1, {}, rescale=False)
+
+    captured = str(excinfo.value)
+    assert "Unsupported op type: 'FakeOp'" in captured
+
+
+# Could be unit or integration?
+@pytest.mark.unit
+def test_check_model_raises_unsupported(dummy_model: ModelProto) -> None:
+    quantizer = ONNXOpQuantizer()
+    quantizer.handlers = {"Dummy": MockHandler()}
+
+    # Remove one node to simulate unsupported ops
+    dummy_model.graph.node.append(helper.make_node("FakeOp", ["a"], ["b"]))
+
+    with pytest.raises(UnsupportedOpError) as excinfo:
+        quantizer.check_model(dummy_model)
+
+    assert "FakeOp" in str(excinfo.value)
+
+
+@pytest.mark.unit
+def test_check_layer_missing_handler(valid_node: NodeProto) -> None:
+    quantizer = ONNXOpQuantizer()
+    with pytest.raises(MissingHandlerError) as exc_info:
+        quantizer.check_layer(valid_node, {})
+
+    assert QuantizationError("").GENERIC_MESSAGE in str(exc_info.value)
+    assert "No quantization handler registered for operator type 'Dummy'." in str(
+        exc_info.value,
+    )
+
+
+@pytest.mark.unit
+def test_check_layer_with_bad_handler(invalid_node: NodeProto) -> None:
+    quantizer = ONNXOpQuantizer()
+    quantizer.handlers = {"Dummy": MockHandler()}
+
+    # This error is created in our mock handler
+    with pytest.raises(ValueError, match="Invalid node parameters"):
+        quantizer.check_layer(invalid_node, {})
+
+
+@pytest.mark.unit
+def test_get_initializer_map_extracts_all() -> None:
+    one_f = 1.0
+    two_f = 2.0
+    count_init = 2
+    tensor1 = helper.make_tensor("a", TensorProto.FLOAT, [1], [one_f])
+    tensor2 = helper.make_tensor("b", TensorProto.FLOAT, [1], [two_f])
+    graph = helper.make_graph([], "g", [], [], initializer=[tensor1, tensor2])
+    model = helper.make_model(graph)
+
+    quantizer = ONNXOpQuantizer()
+    init_map = quantizer.get_initializer_map(model)
+    assert init_map["a"].float_data[0] == one_f
+    assert init_map["b"].float_data[0] == two_f
+
+    # Enhanced: check all properties
+    assert len(init_map) == count_init
+    assert init_map["a"].name == "a"
+    assert init_map["a"].data_type == TensorProto.FLOAT
+    assert init_map["a"].dims == [1]
+    assert init_map["b"].name == "b"
+    assert init_map["b"].data_type == TensorProto.FLOAT
+    assert init_map["b"].dims == [1]
+    # Using numpy_helper for consistency
+    assert onnx.numpy_helper.to_array(init_map["a"])[0] == one_f
+    assert onnx.numpy_helper.to_array(init_map["b"])[0] == two_f
+
+
+@pytest.mark.unit
+def test_check_layer_skips_handler_without_check_supported() -> None:
+    class NoCheckHandler:
+        def quantize(self, *args: tuple, **kwargs: dict[str, Any]) -> None:
+            pass  # no check_supported
+
+    quantizer = ONNXOpQuantizer()
+    quantizer.register("NoCheckOp", NoCheckHandler())
+
+    node = helper.make_node("NoCheckOp", ["x"], ["y"])
+    # Should not raise
+    quantizer.check_layer(node, {})
+
+
+@pytest.mark.unit
+def test_register_overwrites_handler() -> None:
+    quantizer = ONNXOpQuantizer()
+    handler1 = MockHandler()
+    handler2 = MockHandler()
+
+    quantizer.register("Dummy", handler1)
+    quantizer.register("Dummy", handler2)
+
+    assert quantizer.handlers["Dummy"] is handler2
+
+
+@pytest.mark.unit
+def test_check_empty_model() -> None:
+    model = helper.make_model(helper.make_graph([], "empty", [], []))
+    quantizer = ONNXOpQuantizer()
+    # Should not raise
+    quantizer.check_model(model)

python/tests/onnx_quantizer_tests/test_registered_quantizers.py

@@ -0,0 +1,121 @@
+# This file performs very basic integration tests on each registered quantizer
+
+import numpy as np
+import onnx
+import pytest
+from onnx import helper
+
+from python.core.model_processing.onnx_quantizer.layers.base import ScaleConfig
+from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (
+    ONNXOpQuantizer,
+)
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+
+
+@pytest.fixture
+def dummy_graph() -> onnx.GraphProto:
+    return onnx.GraphProto()
+
+
+def mock_initializer_map(input_names: list[str]) -> dict[str, onnx.TensorProto]:
+    rng = np.random.default_rng(TEST_RNG_SEED)
+    return {
+        name: onnx.helper.make_tensor(
+            name=name,
+            data_type=onnx.TensorProto.FLOAT,
+            dims=[2, 2],  # minimal shape
+            vals=rng.random(4, dtype=np.float32).tolist(),
+        )
+        for name in input_names
+    }
+
+
+def get_required_input_names(op_type: str) -> list[str]:
+    try:
+        schema = onnx.defs.get_schema(op_type)
+        return [
+            inp.name or f"input{i}"
+            for i, inp in enumerate(schema.inputs)
+            if inp.option != 1
+        ]  # 1 = optional
+    except Exception:
+        return ["input0"]  # fallback
+
+
+def validate_quantized_node(node_result: onnx.NodeProto, op_type: str) -> None:
+    """Validate a single quantized node."""
+    assert isinstance(node_result, onnx.NodeProto), f"Invalid node type for {op_type}"
+    assert node_result.op_type, f"Missing op_type for {op_type}"
+    assert node_result.output, f"Missing outputs for {op_type}"
+
+    try:
+        # Create a minimal model with custom opset for validation
+        temp_graph = onnx.GraphProto()
+        temp_graph.name = "temp_graph"
+
+        # Add dummy inputs/outputs to satisfy graph requirements
+        for inp in node_result.input:
+            if not any(vi.name == inp for vi in temp_graph.input):
+                temp_graph.input.append(
+                    onnx.helper.make_tensor_value_info(
+                        inp,
+                        onnx.TensorProto.FLOAT,
+                        [1],
+                    ),
+                )
+        for out in node_result.output:
+            if not any(vi.name == out for vi in temp_graph.output):
+                temp_graph.output.append(
+                    onnx.helper.make_tensor_value_info(
+                        out,
+                        onnx.TensorProto.FLOAT,
+                        [1],
+                    ),
+                )
+
+        temp_graph.node.append(node_result)
+        temp_model = onnx.helper.make_model(temp_graph)
+        custom_domain = onnx.helper.make_operatorsetid(
+            domain="ai.onnx.contrib",
+            version=1,
+        )
+        temp_model.opset_import.append(custom_domain)
+        onnx.checker.check_model(temp_model)
+    except onnx.checker.ValidationError as e:
+        pytest.fail(f"ONNX node validation failed for {op_type}: {e}")
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("op_type", list(ONNXOpQuantizer().handlers.keys()))
+def test_registered_quantizer_quantize(
+    op_type: str,
+    dummy_graph: onnx.GraphProto,
+) -> None:
+    quantizer = ONNXOpQuantizer()
+    handler = quantizer.handlers[op_type]
+
+    inputs = get_required_input_names(op_type)
+    dummy_initializer_map = mock_initializer_map(inputs)
+
+    dummy_node = helper.make_node(
+        op_type=op_type,
+        inputs=inputs,
+        outputs=["dummy_output"],
+    )
+
+    result = handler.quantize(
+        node=dummy_node,
+        graph=dummy_graph,
+        scale_config=ScaleConfig(exponent=10, base=2, rescale=True),
+        initializer_map=dummy_initializer_map,
+    )
+    assert result is not None
+
+    # Enhanced assertions: validate result type and structure
+    if isinstance(result, list):
+        assert len(result) > 0, f"Quantize returned empty list for {op_type}"
+        for node_result in result:
+            validate_quantized_node(node_result, op_type)
+    else:
+        assert result.input, f"Missing inputs for {op_type}"
+        validate_quantized_node(result, op_type)

python/tests/onnx_quantizer_tests/testing_helper_functions.py

@@ -0,0 +1,17 @@
+from onnx import ModelProto
+
+
+# Helper to extract input shapes
+def get_input_shapes(model: ModelProto) -> dict:
+    input_shapes = {}
+    for inp in model.graph.input:
+        shape = []
+        for dim in inp.type.tensor_type.shape.dim:
+            if dim.HasField("dim_value"):
+                shape.append(int(dim.dim_value))
+            elif dim.dim_param:
+                shape.append(1)  # Default for dynamic dims
+            else:
+                shape.append(1)
+        input_shapes[inp.name] = shape
+    return input_shapes