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

python/tests/circuit_parent_classes/test_onnx_converter.py

@@ -1,9 +1,11 @@
 # test_converter.py
 import tempfile
+from collections.abc import Generator
 from pathlib import Path
-from typing import Any, Generator
+from typing import Any
 from unittest.mock import MagicMock, patch
 
+import numpy as np
 import onnx
 import onnxruntime as ort
 import pytest
@@ -13,7 +15,7 @@ from onnx import TensorProto, helper
 from python.core.model_processing.converters.onnx_converter import ONNXConverter
 
 
-@pytest.fixture()
+@pytest.fixture
 def temp_model_path(
     tmp_path: Generator[Path, None, None],
 ) -> Generator[Path, Any, None]:
@@ -26,7 +28,7 @@ def temp_model_path(
         model_path.unlink()
 
 
-@pytest.fixture()
+@pytest.fixture
 def temp_quant_model_path(
     tmp_path: Generator[Path, None, None],
 ) -> Generator[Path, Any, None]:
@@ -39,7 +41,7 @@ def temp_quant_model_path(
         model_path.unlink()
 
 
-@pytest.fixture()
+@pytest.fixture
 def converter() -> ONNXConverter:
     conv = ONNXConverter()
     conv.model = MagicMock(name="model")
@@ -47,7 +49,7 @@ def converter() -> ONNXConverter:
     return conv
 
 
-@pytest.mark.unit()
+@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"
@@ -55,7 +57,7 @@ def test_save_model(mock_save: MagicMock, converter: ONNXConverter) -> None:
     mock_save.assert_called_once_with(converter.model, path)
 
 
-@pytest.mark.unit()
+@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")
@@ -68,7 +70,7 @@ def test_load_model(mock_load: MagicMock, converter: ONNXConverter) -> None:
     assert converter.model == fake_model
 
 
-@pytest.mark.unit()
+@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"
@@ -76,7 +78,7 @@ def test_save_quantized_model(mock_save: MagicMock, converter: ONNXConverter) ->
     mock_save.assert_called_once_with(converter.quantized_model, path)
 
 
-@pytest.mark.unit()
+@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")
@@ -108,10 +110,12 @@ def test_load_quantized_model(
     assert converter.quantized_model == fake_model
 
 
-@pytest.mark.unit()
+@pytest.mark.unit
 def test_get_outputs_with_mocked_session(converter: ONNXConverter) -> None:
-    dummy_input = [[1.0]]
+    dummy_input = np.array([[1.0]])  # Use np.ndarray, not list
     dummy_output = [[2.0]]
+    converter.scale_base = 2
+    converter.scale_exponent = 10
 
     mock_sess = MagicMock()
 
@@ -132,7 +136,10 @@ def test_get_outputs_with_mocked_session(converter: ONNXConverter) -> None:
 
     result = converter.get_outputs(dummy_input)
 
-    mock_sess.run.assert_called_once_with(["output"], {"input": dummy_input})
+    # Expect NumPy array to be passed into ort_sess.run()
+    expected_call_inputs = {"input": np.asarray(dummy_input)}
+    mock_sess.run.assert_called_once_with(["output"], expected_call_inputs)
+
     assert result == dummy_output
 
 
@@ -148,7 +155,7 @@ def create_dummy_model() -> onnx.ModelProto:
     return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)])
 
 
-@pytest.mark.integration()
+@pytest.mark.integration
 def test_save_and_load_real_model() -> None:
     converter = ONNXConverter()
     model = create_dummy_model()
@@ -181,10 +188,12 @@ def test_save_and_load_real_model() -> None:
         assert converter.model.graph.node[0].op_type == "Identity"
 
 
-@pytest.mark.integration()
+@pytest.mark.integration
 def test_real_inference_from_onnx() -> None:
     converter = ONNXConverter()
     converter.model = create_dummy_model()
+    converter.scale_base = 2
+    converter.scale_exponent = 10
 
     # Save and load into onnxruntime
     with tempfile.NamedTemporaryFile(suffix=".onnx") as tmp:

python/tests/onnx_quantizer_tests/__init__.py

@@ -0,0 +1 @@
+TEST_RNG_SEED = 2

python/tests/onnx_quantizer_tests/layers/__init__.py

@@ -0,0 +1,13 @@
+from .base import BaseLayerConfigProvider, LayerTestConfig
+from .factory import TestLayerFactory
+
+# Auto-discover and make available all config providers
+# This triggers the discovery process when the package is imported
+_all_configs = TestLayerFactory.get_layer_configs()
+
+# Export the factory and base classes
+__all__ = [
+    "BaseLayerConfigProvider",
+    "LayerTestConfig",
+    "TestLayerFactory",
+]

python/tests/onnx_quantizer_tests/layers/add_config.py

@@ -0,0 +1,102 @@
+import numpy as np
+
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+from python.tests.onnx_quantizer_tests.layers.base import (
+    BaseLayerConfigProvider,
+    LayerTestConfig,
+    LayerTestSpec,
+    e2e_test,
+    edge_case_test,
+    valid_test,
+)
+
+
+class AddConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Add layer"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Add"
+
+    def get_config(self) -> LayerTestConfig:
+        return LayerTestConfig(
+            op_type="Add",
+            valid_inputs=["A", "B"],
+            valid_attributes={},  # Add has no layer-specific attributes
+            required_initializers={},
+            input_shapes={
+                "A": [1, 3, 4, 4],
+                "B": [1, 3, 4, 4],
+            },
+            output_shapes={
+                "add_output": [1, 3, 4, 4],
+            },
+        )
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        return [
+            # --- VALID TESTS ---
+            valid_test("basic")
+            .description("Basic elementwise Add of two same-shaped tensors")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 4, 4])
+            .tags("basic", "elementwise", "add")
+            .build(),
+            valid_test("broadcast_add")
+            .description("Add with Numpy-style broadcasting along spatial dimensions")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 1, 1])
+            .tags("broadcast", "elementwise", "add", "onnx14")
+            .build(),
+            valid_test("initializer_add")
+            .description(
+                "Add where second input (B) is a tensor initializer instead of input",
+            )
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", rng.normal(0, 1, (1, 3, 4, 4)))
+            .tags("initializer", "elementwise", "add", "onnxruntime")
+            .build(),
+            valid_test("scalar_add")
+            .description("Add scalar (initializer) to tensor")
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", np.array([2.0], dtype=np.float32))
+            .tags("scalar", "elementwise", "add")
+            .build(),
+            # --- E2E TESTS ---
+            e2e_test("e2e_add")
+            .description("End-to-end Add test with random inputs")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 4, 4])
+            .override_output_shapes(add_output=[1, 3, 4, 4])
+            .tags("e2e", "add", "2d")
+            .build(),
+            e2e_test("e2e_initializer_add")
+            .description(
+                "Add where second input (B) is a tensor initializer instead of input",
+            )
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", rng.normal(0, 1, (1, 3, 4, 4)))
+            .tags("initializer", "elementwise", "add", "onnxruntime")
+            .build(),
+            e2e_test("e2e_broadcast_add")
+            .description("Add with Numpy-style broadcasting along spatial dimensions")
+            .override_input_shapes(A=[1, 3, 4, 4], B=[1, 3, 1, 1])
+            .tags("broadcast", "elementwise", "add", "onnx14")
+            .build(),
+            e2e_test("e2e_scalar_add")
+            .description("Add scalar (initializer) to tensor")
+            .override_input_shapes(A=[1, 3, 4, 4])
+            .override_initializer("B", np.array([2.0], dtype=np.float32))
+            .tags("scalar", "elementwise", "add")
+            .build(),
+            # # --- EDGE CASES ---
+            edge_case_test("empty_tensor")
+            .description("Add with empty tensor input (zero elements)")
+            .override_input_shapes(A=[0], B=[0])
+            .tags("edge", "empty", "add")
+            .build(),
+            edge_case_test("large_tensor")
+            .description("Large tensor add performance/stress test")
+            .override_input_shapes(A=[1, 64, 256, 256], B=[1, 64, 256, 256])
+            .tags("large", "performance", "add")
+            .skip("Performance test, skipped by default")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers/base.py

@@ -0,0 +1,279 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper
+
+
+class SpecType(Enum):
+    """Types of test specifications that can be run"""
+
+    VALID = "valid"
+    ERROR = "error"
+    EDGE_CASE = "edge_case"
+    E2E = "e2e"
+
+
+@dataclass
+class LayerTestSpec:
+    """Individual test specification that can be applied to a LayerTestConfig"""
+
+    name: str
+    spec_type: SpecType
+    description: str = ""
+
+    # Overrides for the base config
+    attr_overrides: dict[str, Any] = field(default_factory=dict)
+    initializer_overrides: dict[str, np.ndarray] = field(default_factory=dict)
+    input_overrides: list[str] = field(default_factory=list)
+    input_shape_overrides: dict[str, list[int]] = field(default_factory=dict)
+    output_shape_overrides: dict[str, list[int]] = field(default_factory=dict)
+
+    # Error test specific
+    expected_error: type | None = None
+    error_match: str | None = None
+
+    # Custom validation
+    custom_validator: Callable | None = None
+
+    # Test metadata
+    tags: list[str] = field(default_factory=list)
+    skip_reason: str | None = None
+
+    # Omit attributes
+    omit_attrs: list[str] = field(default_factory=list)
+
+    # Remove __post_init__ validation - we'll validate in the builder instead
+
+
+class LayerTestConfig:
+    """Enhanced configuration class for layer-specific test data"""
+
+    def __init__(
+        self: LayerTestConfig,
+        op_type: str,
+        valid_inputs: list[str],
+        valid_attributes: dict[str, Any],
+        required_initializers: dict[str, np.ndarray],
+        input_shapes: dict[str, list[int]] | None = None,
+        output_shapes: dict[str, list[int]] | None = None,
+    ) -> None:
+        self.op_type = op_type
+        self.valid_inputs = valid_inputs
+        self.valid_attributes = valid_attributes
+        self.required_initializers = required_initializers
+        self.input_shapes = input_shapes or {"input": [1, 16, 224, 224]}
+        self.output_shapes = output_shapes or {f"{op_type.lower()}_output": [1, 10]}
+
+    def create_node(
+        self: LayerTestConfig,
+        name_suffix: str = "",
+        **attr_overrides: dict[str, Any],
+    ) -> onnx.NodeProto:
+        """Create a valid node for this layer type"""
+        attrs = {**self.valid_attributes, **attr_overrides}
+        return helper.make_node(
+            self.op_type,
+            inputs=self.valid_inputs,
+            outputs=[f"{self.op_type.lower()}_output{name_suffix}"],
+            name=f"test_{self.op_type.lower()}{name_suffix}",
+            **attrs,
+        )
+
+    def create_initializers(
+        self: LayerTestConfig,
+        **initializer_overrides: dict[str, Any],
+    ) -> dict[str, onnx.TensorProto]:
+        """Create initializer tensors for this layer"""
+        initializers = {}
+        combined_inits = {**self.required_initializers, **initializer_overrides}
+        for name, data in combined_inits.items():
+            # Special handling for shape tensors in Reshape, etc.
+            if name == "shape":
+                tensor = numpy_helper.from_array(data.astype(np.int64), name=name)
+            else:
+                tensor = numpy_helper.from_array(data.astype(np.float32), name=name)
+            initializers[name] = tensor
+        return initializers
+
+    def create_test_model(self, test_spec: LayerTestSpec) -> onnx.ModelProto:
+        """Create a complete model for a specific test case"""
+
+        # Determine node-level inputs.
+        # If dev overrides inputs explicitly,
+        # respect that; otherwise use original valid_inputs.
+        inputs = test_spec.input_overrides or self.valid_inputs
+
+        # Prepare attributes
+        attrs = {**self.valid_attributes, **test_spec.attr_overrides}
+        # Remove omitted attributes if specified
+        attrs = {**self.valid_attributes, **test_spec.attr_overrides}
+        for key in getattr(test_spec, "omit_attrs", []):
+            attrs.pop(key, None)
+
+        # Create initializers (may introduce overrides)
+        initializers = self.create_initializers(**test_spec.initializer_overrides)
+
+        # Apply shape overrides
+        input_shapes = {**self.input_shapes, **test_spec.input_shape_overrides}
+        output_shapes = {**self.output_shapes, **test_spec.output_shape_overrides}
+
+        # ----------------------------------------
+        # REMOVE graph inputs that are also initializers
+        # ----------------------------------------
+        initializer_names = set(initializers.keys())
+
+        # Also remove shapes for initializer inputs
+        for init_name in initializer_names:
+            input_shapes.pop(init_name, None)
+
+        # Create ONNX input value infos ONLY from filtered inputs
+        graph_inputs = [
+            helper.make_tensor_value_info(name, TensorProto.FLOAT, shape)
+            for name, shape in input_shapes.items()
+        ]
+
+        # Outputs stay unchanged
+        graph_outputs = [
+            helper.make_tensor_value_info(name, TensorProto.FLOAT, shape)
+            for name, shape in output_shapes.items()
+        ]
+
+        node = helper.make_node(
+            self.op_type,
+            inputs=inputs,
+            outputs=[f"{self.op_type.lower()}_output"],
+            name=f"test_{self.op_type.lower()}_{test_spec.name}",
+            **attrs,
+        )
+
+        # Build the graph
+        graph = helper.make_graph(
+            nodes=[node],
+            name=f"{self.op_type.lower()}_test_graph_{test_spec.name}",
+            inputs=graph_inputs,
+            outputs=graph_outputs,
+            initializer=list(initializers.values()),
+        )
+
+        return helper.make_model(graph)
+
+
+class TestSpecBuilder:
+    """Builder for creating test specifications"""
+
+    def __init__(self, name: str, spec_type: SpecType) -> None:
+        self._spec = LayerTestSpec(name=name, spec_type=spec_type)
+
+    def description(self, desc: str) -> TestSpecBuilder:
+        self._spec.description = desc
+        return self
+
+    def override_attrs(self, **attrs: dict[str, Any]) -> TestSpecBuilder:
+        self._spec.attr_overrides.update(attrs)
+        return self
+
+    def omit_attrs(self, *attrs: str) -> TestSpecBuilder:
+        self._spec.omit_attrs.extend(attrs)
+        return self
+
+    def override_initializer(self, name: str, data: np.ndarray) -> TestSpecBuilder:
+        self._spec.initializer_overrides[name] = data
+        return self
+
+    def override_inputs(self, *inputs: str) -> TestSpecBuilder:
+        self._spec.input_overrides = list(inputs)
+        return self
+
+    def override_input_shapes(self, **shapes: dict[str, list[int]]) -> TestSpecBuilder:
+        self._spec.input_shape_overrides.update(shapes)
+        return self
+
+    def override_output_shapes(self, **shapes: dict[str, list[int]]) -> TestSpecBuilder:
+        self._spec.output_shape_overrides.update(shapes)
+        return self
+
+    def expects_error(
+        self,
+        error_type: type,
+        match: str | None = None,
+    ) -> TestSpecBuilder:
+        if self._spec.spec_type != SpecType.ERROR:
+            msg = "expects_error can only be used with ERROR spec type"
+            raise ValueError(msg)
+        self._spec.expected_error = error_type
+        self._spec.error_match = match
+        return self
+
+    def tags(self, *tags: str) -> TestSpecBuilder:
+        self._spec.tags.extend(tags)
+        return self
+
+    def skip(self, reason: str) -> TestSpecBuilder:
+        self._spec.skip_reason = reason
+        return self
+
+    def build(self) -> LayerTestSpec:
+        # Validate before building
+        if self._spec.spec_type == SpecType.ERROR and not self._spec.expected_error:
+            msg = (
+                f"Error test {self._spec.name} must"
+                " specify expected_error using .expects_error()"
+            )
+            raise ValueError(msg)
+        return self._spec
+
+
+# Convenience functions
+def valid_test(name: str) -> TestSpecBuilder:
+    return TestSpecBuilder(name, SpecType.VALID)
+
+
+def error_test(name: str) -> TestSpecBuilder:
+    return TestSpecBuilder(name, SpecType.ERROR)
+
+
+def edge_case_test(name: str) -> TestSpecBuilder:
+    return TestSpecBuilder(name, SpecType.EDGE_CASE)
+
+
+def e2e_test(name: str) -> TestSpecBuilder:
+    return TestSpecBuilder(name, SpecType.E2E)
+
+
+class BaseLayerConfigProvider(ABC):
+    """Abstract base class for layer config providers"""
+
+    @abstractmethod
+    def get_config(self) -> LayerTestConfig:
+        """Return the base configuration for this layer"""
+
+    @property
+    @abstractmethod
+    def layer_name(self) -> str:
+        """Return the layer name/op_type"""
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        """Return test specifications for this layer (override for custom tests)"""
+        return []
+
+    def get_valid_test_specs(self) -> list[LayerTestSpec]:
+        """Get only valid test specifications"""
+        return [
+            spec for spec in self.get_test_specs() if spec.spec_type == SpecType.VALID
+        ]
+
+    def get_error_test_specs(self) -> list[LayerTestSpec]:
+        """Get only error test specifications"""
+        return [
+            spec for spec in self.get_test_specs() if spec.spec_type == SpecType.ERROR
+        ]