JSTprove 1.0.0__py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

python/tests/circuit_parent_classes/test_onnx_converter.py

@@ -0,0 +1,201 @@
+# test_converter.py
+import tempfile
+from pathlib import Path
+from typing import Any, Generator
+from unittest.mock import MagicMock, patch
+
+import onnx
+import onnxruntime as ort
+import pytest
+import torch
+from onnx import TensorProto, helper
+
+from python.core.model_processing.converters.onnx_converter import ONNXConverter
+
+
+@pytest.fixture()
+def temp_model_path(
+    tmp_path: Generator[Path, None, None],
+) -> Generator[Path, Any, None]:
+    model_path = tmp_path / "temp_model.onnx"
+    # Give it to the test
+    yield model_path
+
+    # After the test is done, remove it
+    if Path.exists(model_path):
+        model_path.unlink()
+
+
+@pytest.fixture()
+def temp_quant_model_path(
+    tmp_path: Generator[Path, None, None],
+) -> Generator[Path, Any, None]:
+    model_path = tmp_path / "temp_quantized_model.onnx"
+    # Give it to the test
+    yield model_path
+
+    # After the test is done, remove it
+    if Path.exists(model_path):
+        model_path.unlink()
+
+
+@pytest.fixture()
+def converter() -> ONNXConverter:
+    conv = ONNXConverter()
+    conv.model = MagicMock(name="model")
+    conv.quantized_model = MagicMock(name="quantized_model")
+    return conv
+
+
+@pytest.mark.unit()
+@patch("python.core.model_processing.converters.onnx_converter.onnx.save")
+def test_save_model(mock_save: MagicMock, converter: ONNXConverter) -> None:
+    path = "model.onnx"
+    converter.save_model(path)
+    mock_save.assert_called_once_with(converter.model, path)
+
+
+@pytest.mark.unit()
+@patch("python.core.model_processing.converters.onnx_converter.onnx.load")
+def test_load_model(mock_load: MagicMock, converter: ONNXConverter) -> None:
+    fake_model = MagicMock(name="onnx_model")
+    mock_load.return_value = fake_model
+
+    path = "model.onnx"
+    converter.load_model(path)
+
+    mock_load.assert_called_once_with(path)
+    assert converter.model == fake_model
+
+
+@pytest.mark.unit()
+@patch("python.core.model_processing.converters.onnx_converter.onnx.save")
+def test_save_quantized_model(mock_save: MagicMock, converter: ONNXConverter) -> None:
+    path = "quantized_model.onnx"
+    converter.save_quantized_model(path)
+    mock_save.assert_called_once_with(converter.quantized_model, path)
+
+
+@pytest.mark.unit()
+@patch("python.core.model_processing.converters.onnx_converter.Path.exists")
+@patch("python.core.model_processing.converters.onnx_converter.SessionOptions")
+@patch("python.core.model_processing.converters.onnx_converter.InferenceSession")
+@patch("python.core.model_processing.converters.onnx_converter.onnx.load")
+def test_load_quantized_model(
+    mock_load: MagicMock,
+    mock_ort_sess: MagicMock,
+    mock_session_opts: MagicMock,
+    mock_exists: MagicMock,
+    converter: ONNXConverter,
+) -> None:
+
+    fake_model = MagicMock(name="onnx_model")
+    mock_load.return_value = fake_model
+    mock_exists.return_value = True  # Mock os.path.exists to return True
+
+    mock_opts_instance = MagicMock(name="session_options")
+    mock_session_opts.return_value = mock_opts_instance
+
+    path = "quantized_model.onnx"
+    converter.load_quantized_model(path)
+
+    mock_load.assert_called_once_with(path)
+    mock_ort_sess.assert_called_once_with(
+        path,
+        mock_opts_instance,
+        providers=["CPUExecutionProvider"],
+    )
+    assert converter.quantized_model == fake_model
+
+
+@pytest.mark.unit()
+def test_get_outputs_with_mocked_session(converter: ONNXConverter) -> None:
+    dummy_input = [[1.0]]
+    dummy_output = [[2.0]]
+
+    mock_sess = MagicMock()
+
+    # Mock .get_inputs()[0].name => "input"
+    mock_input = MagicMock()
+    mock_input.name = "input"
+    mock_sess.get_inputs.return_value = [mock_input]
+
+    # Mock .get_outputs()[0].name => "output"
+    mock_output = MagicMock()
+    mock_output.name = "output"
+    mock_sess.get_outputs.return_value = [mock_output]
+
+    # Mock .run() output
+    mock_sess.run.return_value = dummy_output
+
+    converter.ort_sess = mock_sess
+
+    result = converter.get_outputs(dummy_input)
+
+    mock_sess.run.assert_called_once_with(["output"], {"input": dummy_input})
+    assert result == dummy_output
+
+
+# Integration test
+
+
+def create_dummy_model() -> onnx.ModelProto:
+    input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])
+    output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])
+    node = helper.make_node("Identity", inputs=["input"], outputs=["output"])
+    graph = helper.make_graph([node], "test-graph", [input_tensor], [output_tensor])
+
+    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
+
+
+@pytest.mark.integration()
+def test_save_and_load_real_model() -> None:
+    converter = ONNXConverter()
+    model = create_dummy_model()
+    converter.model = model
+    converter.quantized_model = model
+
+    with tempfile.NamedTemporaryFile(suffix=".onnx") as tmp:
+        # Save model
+        converter.save_model(tmp.name)
+
+        # Load model
+        converter.load_model(tmp.name)
+
+        # Validate loaded model
+        assert isinstance(converter.model, onnx.ModelProto)
+        assert converter.model.graph.name == model.graph.name
+        assert len(converter.model.graph.node) == 1
+        assert converter.model.graph.node[0].op_type == "Identity"
+
+        # Save model
+        converter.save_quantized_model(tmp.name)
+
+        # Load model
+        converter.load_quantized_model(tmp.name)
+
+        # Validate loaded model
+        assert isinstance(converter.model, onnx.ModelProto)
+        assert converter.model.graph.name == model.graph.name
+        assert len(converter.model.graph.node) == 1
+        assert converter.model.graph.node[0].op_type == "Identity"
+
+
+@pytest.mark.integration()
+def test_real_inference_from_onnx() -> None:
+    converter = ONNXConverter()
+    converter.model = create_dummy_model()
+
+    # Save and load into onnxruntime
+    with tempfile.NamedTemporaryFile(suffix=".onnx") as tmp:
+        onnx.save(converter.model, tmp.name)
+        converter.ort_sess = ort.InferenceSession(
+            tmp.name,
+            providers=["CPUExecutionProvider"],
+        )
+
+        dummy_input = torch.tensor([1.0], dtype=torch.float32).numpy()
+        result = converter.get_outputs(dummy_input)
+
+        assert isinstance(result, list)
+        print(result)  # Identity op should return input

python/tests/circuit_parent_classes/test_ort_custom_layers.py

@@ -0,0 +1,116 @@
+import pytest
+import numpy as np
+import torch
+import onnx
+
+from onnx import TensorProto, shape_inference, helper, numpy_helper
+
+from python.core.model_processing.converters.onnx_converter import ONNXConverter
+from python.core.model_processing.onnx_custom_ops.onnx_helpers import extract_shape_dict
+from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import ONNXOpQuantizer
+
+from onnxruntime import InferenceSession, SessionOptions
+from onnxruntime_extensions import get_library_path, OrtPyFunction
+from python.core.model_processing.onnx_custom_ops import conv
+
+from python.core.model_processing.onnx_custom_ops.conv import int64_conv
+from python.core.model_processing.onnx_custom_ops.gemm import int64_gemm7
+
+
+@pytest.fixture
+def tiny_conv_model_path(tmp_path):
+    # Create input and output tensor info
+    input_tensor = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 1, 4, 4])
+    output_tensor = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 1, 2, 2])
+
+    # Kernel weights (3x3 ones)
+    W_init = helper.make_tensor(
+        name='W',
+        data_type=TensorProto.FLOAT,
+        dims=[1, 1, 3, 3],
+        vals=np.ones((1 * 1 * 3 * 3), dtype=np.float32).tolist()
+    )
+    Z_init = helper.make_tensor(
+        name='Z',
+        data_type=TensorProto.FLOAT,
+        dims=[1],
+        vals=np.ones(( 1), dtype=np.float32).tolist()
+    )
+
+    # Conv node with no padding, stride 1
+    conv_node = helper.make_node(
+        'Conv',
+        inputs=['X', 'W', 'Z'],
+        outputs=['Y'],
+        kernel_shape=[3, 3],
+        pads=[0, 0, 0, 0],
+        strides=[1, 1],
+        dilations = [1,1]
+    )
+
+    # Build graph and model
+    graph = helper.make_graph(
+        nodes=[conv_node],
+        name='TinyConvGraph',
+        inputs=[input_tensor],
+        outputs=[output_tensor],
+        initializer=[W_init, Z_init]
+    )
+
+    model = helper.make_model(graph, producer_name='tiny-conv-example')
+
+    # Save to a temporary file
+    model_path = tmp_path / "tiny_conv.onnx"
+    onnx.save(model, str(model_path))
+
+    return str(model_path)
+
+@pytest.mark.integration
+def test_tiny_conv(tiny_conv_model_path):
+    path = tiny_conv_model_path
+
+    converter = ONNXConverter()
+
+    X_input = np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4)
+    id_count = 0
+    model = onnx.load(path)
+    # Fix, can remove this next line 
+    onnx.checker.check_model(model)
+
+    # Check the model and print Y"s shape information
+    onnx.checker.check_model(model)
+    print(f"Before shape inference, the shape info of Y is:\n{model.graph.value_info}")
+
+    # Apply shape inference on the model
+    inferred_model = shape_inference.infer_shapes(model)
+
+    # Check the model and print Y"s shape information
+    onnx.checker.check_model(inferred_model)
+    # print(f"After shape inference, the shape info of Y is:\n{inferred_model.graph.value_info}")
+    
+
+    domain_to_version = {opset.domain: opset.version for opset in model.opset_import}
+    
+    inferred_model = shape_inference.infer_shapes(model)
+    output_name_to_shape = extract_shape_dict(inferred_model)
+    id_count = 0
+
+    new_model = converter.quantize_model(model, 2, 21)
+    custom_domain = onnx.helper.make_operatorsetid(domain="ai.onnx.contrib", version=1)
+    new_model.opset_import.append(custom_domain)
+    onnx.checker.check_model(new_model)
+
+    with open("model.onnx", "wb") as f:
+        f.write(new_model.SerializeToString())
+
+    model = onnx.load("model.onnx")
+    onnx.checker.check_model(model)  # This throws a descriptive error
+
+    inputs = np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4)
+    outputs_true = converter.run_model_onnx_runtime(path, inputs)
+
+    outputs_quant = converter.run_model_onnx_runtime("model.onnx", inputs)
+    true = torch.tensor(np.array(outputs_true), dtype=torch.float32)
+    quant = torch.tensor(np.array(outputs_quant), dtype=torch.float32) / (2**21)
+
+    assert torch.allclose(true, quant, rtol=1e-3, atol=1e-5), "Outputs do not match"