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/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

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)