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

@@ -0,0 +1,61 @@
+import numpy as np
+
+from python.tests.onnx_quantizer_tests.layers.base import (
+    e2e_test,
+    valid_test,
+)
+from python.tests.onnx_quantizer_tests.layers.factory import (
+    BaseLayerConfigProvider,
+    LayerTestConfig,
+)
+
+
+class ReshapeConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Reshape layers"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Reshape"
+
+    def get_config(self) -> LayerTestConfig:
+        return LayerTestConfig(
+            op_type="Reshape",
+            valid_inputs=["input", "shape"],
+            valid_attributes={},
+            required_initializers={"shape": np.array([1, -1])},
+        )
+
+    def get_test_specs(self) -> list:
+        return [
+            # --- VALID TESTS ---
+            valid_test("basic")
+            .description("Basic Reshape from (1,2,3,4) to (1,24)")
+            .tags("basic", "reshape")
+            .build(),
+            valid_test("reshape_expand_dims")
+            .description("Reshape expanding dimensions (1,24) → (1,3,8)")
+            .override_input_shapes(input=[1, 24])
+            .tags("reshape", "expand")
+            .build(),
+            valid_test("reshape_flatten")
+            .description("Reshape to flatten spatial dimensions (1,3,4,4) → (1,48)")
+            .override_input_shapes(input=[1, 24])
+            .override_initializer("shape", np.array([1, 3, -1]))
+            .tags("reshape", "flatten")
+            .build(),
+            e2e_test("e2e_basic")
+            .description("End-to-end test for Reshape layer")
+            .override_input_shapes(input=[1, 2, 3, 4])
+            .override_output_shapes(reshape_output=[1, 24])
+            .override_initializer("shape", np.array([1, -1]))
+            .tags("e2e", "reshape")
+            .build(),
+            # --- EDGE CASE / SKIPPED TEST ---
+            valid_test("large_input")
+            .description("Large reshape performance test")
+            .override_input_shapes(input=[1, 3, 256, 256])
+            .override_initializer("shape", np.array([1, -1]))
+            .tags("large", "performance", "reshape")
+            # .skip("Performance test, skipped by default")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers_tests/base_test.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import pytest
+from onnx import TensorProto, helper
+
+if TYPE_CHECKING:
+    from onnx import ModelProto
+
+    from python.tests.onnx_quantizer_tests.layers.base import (
+        LayerTestConfig,
+        LayerTestSpec,
+    )
+from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (
+    ONNXOpQuantizer,
+)
+from python.tests.onnx_quantizer_tests.layers.factory import TestLayerFactory
+
+
+class BaseQuantizerTest:
+    """Base test utilities for ONNX quantizer tests."""
+
+    __test__ = False  # Prevent pytest from collecting this class directly
+
+    _validation_failed_cases: ClassVar[set[str]] = set()
+
+    @pytest.fixture
+    def quantizer(self) -> ONNXOpQuantizer:
+        return ONNXOpQuantizer()
+
+    @pytest.fixture
+    def layer_configs(self) -> dict[str, LayerTestConfig]:
+        return TestLayerFactory.get_layer_configs()
+
+    @staticmethod
+    def _generate_test_id(
+        test_case_tuple: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> str:
+        try:
+            layer_name, _, test_spec = test_case_tuple
+        except Exception:
+            return str(test_case_tuple)
+        else:
+            return f"{layer_name}_{test_spec.name}"
+
+    @classmethod
+    def _check_validation_dependency(
+        cls: BaseQuantizerTest,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        layer_name, _, test_spec = test_case_data
+        test_case_id = f"{layer_name}_{test_spec.name}"
+        if test_case_id in cls._validation_failed_cases:
+            pytest.skip(f"Skipping because ONNX validation failed for {test_case_id}")
+
+    @staticmethod
+    def create_model_with_layers(
+        layer_types: list[str],
+        layer_configs: dict[str, LayerTestConfig],
+    ) -> ModelProto:
+        """Create a model composed of several layers."""
+        nodes, all_initializers = [], {}
+
+        for i, layer_type in enumerate(layer_types):
+            config = layer_configs[layer_type]
+            node = config.create_node(name_suffix=f"_{i}")
+            if i > 0:
+                prev_output = f"{layer_types[i-1].lower()}_output_{i-1}"
+                if node.input:
+                    node.input[0] = prev_output
+            nodes.append(node)
+            all_initializers.update(config.create_initializers())
+
+        graph = helper.make_graph(
+            nodes,
+            "test_graph",
+            [
+                helper.make_tensor_value_info(
+                    "input",
+                    TensorProto.FLOAT,
+                    [1, 16, 224, 224],
+                ),
+            ],
+            [
+                helper.make_tensor_value_info(
+                    f"{layer_types[-1].lower()}_output_{len(layer_types)-1}",
+                    TensorProto.FLOAT,
+                    [1, 10],
+                ),
+            ],
+            initializer=list(all_initializers.values()),
+        )
+        return helper.make_model(graph)

python/tests/onnx_quantizer_tests/layers_tests/test_check_model.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from onnx import TensorProto, helper
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    InvalidParamError,
+    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 TestCheckModel(BaseQuantizerTest):
+    """Tests for ONNX model checking."""
+
+    __test__ = True
+
+    @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_check_model_individual_valid_cases(
+        self: TestCheckModel,
+        quantizer: ONNXOpQuantizer,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        """Test each individual valid test case"""
+        layer_name, config, test_spec = test_case_data
+
+        # Skips if layer is not a valid onnx layer
+        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 model from layer specs
+        model = config.create_test_model(test_spec)
+
+        try:
+            quantizer.check_model(model)
+        except (InvalidParamError, UnsupportedOpError) as e:
+            pytest.fail(f"Model check failed for {layer_name}.{test_spec.name}: {e}")
+        except Exception as e:
+            pytest.fail(f"Model check failed for {layer_name}.{test_spec.name}: {e}")
+
+    @pytest.mark.unit
+    def test_check_model_unsupported_layer_fails(
+        self: TestCheckModel,
+        quantizer: ONNXOpQuantizer,
+    ) -> None:
+        """Test that models with unsupported layers fail validation"""
+        # Create model with unsupported operation
+        unsupported_node = helper.make_node(
+            "UnsupportedOp",
+            inputs=["input"],
+            outputs=["output"],
+            name="unsupported",
+        )
+
+        graph = helper.make_graph(
+            [unsupported_node],
+            "test_graph",
+            [
+                helper.make_tensor_value_info(
+                    "input",
+                    TensorProto.FLOAT,
+                    [1, 16, 224, 224],
+                ),
+            ],
+            [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 10])],
+        )
+
+        model = helper.make_model(graph)
+
+        with pytest.raises(UnsupportedOpError):
+            quantizer.check_model(model)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "layer_combination",
+        [
+            ["Conv", "Relu"],
+            ["Conv", "Relu", "MaxPool"],
+            ["Gemm", "Relu"],
+            ["Conv", "Reshape", "Gemm"],
+            ["Conv", "Flatten", "Gemm"],
+        ],
+    )
+    def test_check_model_multi_layer_passes(
+        self: TestCheckModel,
+        quantizer: ONNXOpQuantizer,
+        layer_configs: dict[str, LayerTestConfig],
+        layer_combination: list[str],
+    ) -> None:
+        """Test that models with multiple supported layers pass validation"""
+        model = self.create_model_with_layers(layer_combination, layer_configs)
+        # Should not raise any exception
+        quantizer.check_model(model)

python/tests/onnx_quantizer_tests/layers_tests/test_e2e.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from python.core.circuit_models.generic_onnx import GenericModelONNX
+from python.core.utils.helper_functions import CircuitExecutionConfig, RunType
+from python.tests.onnx_quantizer_tests.layers.base import SpecType
+from python.tests.onnx_quantizer_tests.layers.factory import TestLayerFactory
+from python.tests.onnx_quantizer_tests.layers_tests.base_test import BaseQuantizerTest
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from pathlib import Path
+
+    from python.tests.onnx_quantizer_tests.layers.base import (
+        LayerTestConfig,
+        LayerTestSpec,
+    )
+
+
+class TestE2EQuantizer(BaseQuantizerTest):
+    """End-to-end tests for ONNX quantizer layers."""
+
+    __test__ = True
+
+    @pytest.fixture
+    def temp_quantized_model(self, tmp_path: Path) -> Generator[Path, None, None]:
+        """Temporary path for quantized model."""
+        path = tmp_path / "quantized_model.onnx"
+        yield path
+        if path.exists():
+            path.unlink()
+
+    @pytest.fixture
+    def temp_circuit_path(self, tmp_path: Path) -> Generator[Path, None, None]:
+        """Temporary path for circuit file."""
+        path = tmp_path / "circuit.txt"
+        yield path
+        if path.exists():
+            path.unlink()
+
+    @pytest.fixture
+    def temp_witness_file(self, tmp_path: Path) -> Generator[Path, None, None]:
+        """Temporary path for witness file."""
+        path = tmp_path / "witness.bin"
+        yield path
+        if path.exists():
+            path.unlink()
+
+    @pytest.fixture
+    def temp_input_file(self, tmp_path: Path) -> Generator[Path, None, None]:
+        """Temporary path for input JSON file."""
+        path = tmp_path / "input.json"
+        yield path
+        if path.exists():
+            path.unlink()
+
+    @pytest.fixture
+    def temp_output_file(self, tmp_path: Path) -> Generator[Path, None, None]:
+        """Temporary path for output JSON file."""
+        path = tmp_path / "output.json"
+        yield path
+        if path.exists():
+            path.unlink()
+
+    @pytest.fixture
+    def temp_proof_file(self, tmp_path: Path) -> Generator[Path, None, None]:
+        """Temporary path for proof file."""
+        path = tmp_path / "proof.bin"
+        yield path
+        if path.exists():
+            path.unlink()
+
+    @pytest.mark.e2e
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.E2E),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_e2e_quantize_compile_witness_prove_verify(
+        self,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+        temp_quantized_model: Path,
+        temp_circuit_path: Path,
+        temp_witness_file: Path,
+        temp_input_file: Path,
+        temp_output_file: Path,
+        temp_proof_file: Path,
+        capsys: pytest.CaptureFixture[str],
+    ) -> None:
+        """Test end-to-end flow: quantize model, compile circuit,
+        generate witness, prove, and verify."""
+        layer_name, config, test_spec = test_case_data
+
+        # 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}")
+
+        if layer_name == "Constant":
+            pytest.skip(f"No e2e test for layer {layer_name} yet")
+
+        # Create original model
+        original_model = config.create_test_model(test_spec)
+
+        # Save quantized model to temp location
+        import onnx  # noqa: PLC0415
+
+        onnx.save(original_model, str(temp_quantized_model))
+
+        # Create GenericONNX model instance
+        model = GenericModelONNX(model_name=str(temp_quantized_model))
+
+        # Step 1: Compile circuit
+        model.base_testing(
+            CircuitExecutionConfig(
+                run_type=RunType.COMPILE_CIRCUIT,
+                dev_mode=False,
+                circuit_path=str(temp_circuit_path),
+            ),
+        )
+
+        # Verify circuit file exists
+        assert (
+            temp_circuit_path.exists()
+        ), f"Circuit file not created for {layer_name}.{test_spec.name}"
+
+        # Step 2: Generate witness
+        model.base_testing(
+            CircuitExecutionConfig(
+                run_type=RunType.GEN_WITNESS,
+                dev_mode=False,
+                witness_file=temp_witness_file,
+                circuit_path=str(temp_circuit_path),
+                input_file=temp_input_file,
+                output_file=temp_output_file,
+                write_json=True,
+            ),
+        )
+        # Verify witness and output files exist
+        assert (
+            temp_witness_file.exists()
+        ), f"Witness file not generated for {layer_name}.{test_spec.name}"
+        assert (
+            temp_output_file.exists()
+        ), f"Output file not generated for {layer_name}.{test_spec.name}"
+
+        # Step 3: Prove
+        model.base_testing(
+            CircuitExecutionConfig(
+                run_type=RunType.PROVE_WITNESS,
+                dev_mode=False,
+                witness_file=temp_witness_file,
+                circuit_path=str(temp_circuit_path),
+                input_file=temp_input_file,
+                output_file=temp_output_file,
+                proof_file=temp_proof_file,
+            ),
+        )
+
+        # Verify proof file exists
+        assert (
+            temp_proof_file.exists()
+        ), f"Proof file not generated for {layer_name}.{test_spec.name}"
+
+        # Step 4: Verify
+        model.base_testing(
+            CircuitExecutionConfig(
+                run_type=RunType.GEN_VERIFY,
+                dev_mode=False,
+                witness_file=temp_witness_file,
+                circuit_path=str(temp_circuit_path),
+                input_file=temp_input_file,
+                output_file=temp_output_file,
+                proof_file=temp_proof_file,
+            ),
+        )
+
+        # Capture output and check for success indicators
+        captured = capsys.readouterr()
+        stdout = captured.out
+        stderr = captured.err
+
+        assert stderr == "", "Errors occurred during e2e test for "
+        f"{layer_name}.{test_spec.name}: {stderr}"
+
+        # Check for expected success messages (similar to circuit e2e tests)
+        assert (
+            "Witness Generated" in stdout
+        ), f"Witness generation failed for {layer_name}.{test_spec.name}"
+        assert "Proved" in stdout, f"Proving failed for {layer_name}.{test_spec.name}"
+        assert (
+            "Verified" in stdout
+        ), f"Verification failed for {layer_name}.{test_spec.name}"

python/tests/onnx_quantizer_tests/layers_tests/test_error_cases.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+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 TestErrorCases(BaseQuantizerTest):
+    """Tests for ONNX model checking."""
+
+    __test__ = True
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "test_case_data",
+        TestLayerFactory.get_test_cases_by_type(SpecType.ERROR),  # type: ignore[arg-type]
+        ids=BaseQuantizerTest._generate_test_id,
+    )
+    def test_check_model_individual_error_cases(
+        self: TestErrorCases,
+        quantizer: ONNXOpQuantizer,
+        test_case_data: tuple[str, LayerTestConfig, LayerTestSpec],
+    ) -> None:
+        """Test each individual error test case"""
+        layer_name, config, test_spec = test_case_data
+
+        # Skips if layer is not a valid onnx layer
+        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 model from layer specs
+        model = config.create_test_model(test_spec)
+
+        # Ensures that expected test is in fact raised
+        with pytest.raises(test_spec.expected_error) as exc:
+            quantizer.check_model(model)
+
+        # Ensures the error message is as expected
+        if isinstance(test_spec.error_match, list):
+            for e in test_spec.error_match:
+                assert e in str(exc.value)
+        else:
+            assert test_spec.error_match in str(exc.value)