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/layers/batchnorm_config.py

@@ -0,0 +1,190 @@
+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,
+    valid_test,
+)
+
+
+class BatchNormConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for BatchNorm (ONNX BatchNormalization op)"""
+
+    @property
+    def layer_name(self) -> str:
+        return "BatchNormalization"
+
+    def get_config(self) -> LayerTestConfig:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+
+        # default shapes: N x C x H x W
+        default_input_shape = [1, 3, 4, 4]
+        c = default_input_shape[1]
+
+        # typical required initializers (scale, bias, mean, var) are length C
+        return LayerTestConfig(
+            op_type="BatchNormalization",
+            valid_inputs=["X", "scale", "B", "input_mean", "input_var"],
+            valid_attributes={
+                "epsilon": 1e-5,
+                "momentum": 0.9,
+                "training_mode": 0,
+            },
+            required_initializers={
+                # Defaults are stored as numpy arrays with shape (C,)
+                "scale": rng.normal(1.0, 0.5, c).astype(np.float32),
+                "B": rng.normal(0.0, 0.5, c).astype(np.float32),
+                "input_mean": rng.normal(0.0, 1.0, c).astype(np.float32),
+                "input_var": np.abs(rng.normal(1.0, 0.5, c)).astype(np.float32),
+            },
+            input_shapes={"X": default_input_shape},
+            output_shapes={"batchnormalization_output": default_input_shape},
+        )
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        c = 3
+
+        return [
+            # Basic valid tests
+            valid_test("basic_inference")
+            .description("Basic BatchNormalization inference: standard shapes")
+            .tags("basic", "inference", "batchnorm")
+            .build(),
+            valid_test("different_input_shape")
+            .description("Inference with different spatial dims")
+            .override_input_shapes(X=[2, c, 8, 8])
+            .override_output_shapes(batchnormalization_output=[2, c, 8, 8])
+            .tags("inference", "spatial")
+            .build(),
+            valid_test("epsilon_variation")
+            .description("Inference with larger epsilon for numerical stability")
+            .override_attrs(epsilon=1e-3)
+            .tags("epsilon")
+            .build(),
+            valid_test("momentum_variation")
+            .description(
+                "Inference with non-default momentum (has no effect in inference mode)",
+            )
+            .override_attrs(momentum=0.5)
+            .tags("momentum")
+            .build(),
+            valid_test("zero_mean_input")
+            .description("Input with zero mean")
+            .override_initializer("input_mean", np.zeros((c,), dtype=np.float32))
+            .tags("edge", "zero_mean")
+            .build(),
+            # Scalar / broadcast style tests
+            valid_test("per_channel_zero_variance")
+            .description(
+                "Edge case: very small variance values (clamped by epsilon), inference",
+            )
+            .override_initializer("input_var", np.full((c,), 1e-8, dtype=np.float32))
+            .override_attrs(epsilon=1e-5)
+            .tags("edge", "small_variance")
+            .build(),
+            # E2E tests that set explicit initializer values
+            e2e_test("e2e_inference")
+            .description("E2E inference test with explicit initializers")
+            .override_input_shapes(X=[1, c, 2, 2])
+            .override_output_shapes(batchnormalization_output=[1, c, 2, 2])
+            .override_initializer("scale", rng.normal(1.0, 0.1, c).astype(np.float32))
+            .override_initializer("B", rng.normal(0.0, 0.1, c).astype(np.float32))
+            .override_initializer(
+                "input_mean",
+                rng.normal(0.0, 0.1, c).astype(np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.abs(rng.normal(0.5, 0.2, c)).astype(np.float32),
+            )
+            .tags("e2e", "inference")
+            .build(),
+            e2e_test("e2e_inference_small_2x2")
+            .description("E2E inference with small 2x2 spatial input")
+            .override_input_shapes(X=[1, 3, 2, 2])
+            .override_output_shapes(batchnormalization_output=[1, 3, 2, 2])
+            .override_initializer("scale", np.array([1.0, 0.9, 1.1], dtype=np.float32))
+            .override_initializer("B", np.array([0.0, 0.1, -0.1], dtype=np.float32))
+            .override_initializer(
+                "input_mean",
+                np.array([0.5, -0.5, 0.0], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([0.25, 0.5, 0.1], dtype=np.float32),
+            )
+            .tags("e2e", "small", "2x2")
+            .build(),
+            e2e_test("e2e_inference_wide_input")
+            .description("E2E inference with wider input shape (C=4, H=2, W=8)")
+            .override_input_shapes(X=[2, 4, 2, 8])
+            .override_output_shapes(batchnormalization_output=[2, 4, 2, 8])
+            .override_initializer(
+                "scale",
+                np.array([1.0, 0.8, 1.2, 0.9], dtype=np.float32),
+            )
+            .override_initializer(
+                "B",
+                np.array([0.0, 0.1, -0.1, 0.05], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_mean",
+                np.array([0.0, 0.5, -0.5, 0.2], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([1.0, 0.5, 0.25, 0.1], dtype=np.float32),
+            )
+            .tags("e2e", "wide", "C4")
+            .build(),
+            e2e_test("e2e_inference_batch2_channels3")
+            .description("E2E inference with batch size 2 and 3 channels")
+            .override_input_shapes(X=[2, 3, 4, 4])
+            .override_output_shapes(batchnormalization_output=[2, 3, 4, 4])
+            .override_initializer("scale", np.array([0.5, 1.0, 1.5], dtype=np.float32))
+            .override_initializer("B", np.array([0.0, 0.0, 0.0], dtype=np.float32))
+            .override_initializer(
+                "input_mean",
+                np.array([-0.5, 0.0, 0.5], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([0.2, 0.5, 0.8], dtype=np.float32),
+            )
+            .tags("e2e", "batch2", "C3")
+            .build(),
+            e2e_test("e2e_inference_high_epsilon")
+            .description("E2E inference with high epsilon for numerical stability")
+            .override_input_shapes(X=[1, 2, 4, 4])
+            .override_output_shapes(batchnormalization_output=[1, 2, 4, 4])
+            .override_initializer("scale", np.array([1.0, 1.0], dtype=np.float32))
+            .override_initializer("B", np.array([0.1, -0.1], dtype=np.float32))
+            .override_initializer("input_mean", np.array([0.0, 0.5], dtype=np.float32))
+            .override_initializer(
+                "input_var",
+                np.array([0.0, 0.0], dtype=np.float32),
+            )  # tiny variance
+            .override_attrs(epsilon=1e-2)
+            .tags("e2e", "high_epsilon", "numerical_stability")
+            .build(),
+            e2e_test("e2e_inference_non_square")
+            .description("E2E inference with non-square spatial dimensions")
+            .override_input_shapes(X=[1, 3, 2, 5])
+            .override_output_shapes(batchnormalization_output=[1, 3, 2, 5])
+            .override_initializer("scale", np.array([1.0, 0.9, 1.1], dtype=np.float32))
+            .override_initializer("B", np.array([0.0, 0.1, -0.1], dtype=np.float32))
+            .override_initializer(
+                "input_mean",
+                np.array([0.1, -0.1, 0.0], dtype=np.float32),
+            )
+            .override_initializer(
+                "input_var",
+                np.array([0.5, 0.25, 0.75], dtype=np.float32),
+            )
+            .tags("e2e", "non_square", "C3")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers/constant_config.py

@@ -0,0 +1,39 @@
+import numpy as np
+from onnx import numpy_helper
+
+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 ConstantConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Constant layers"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Constant"
+
+    def get_config(self) -> LayerTestConfig:
+        return LayerTestConfig(
+            op_type="Constant",
+            valid_inputs=[],
+            valid_attributes={
+                "value": numpy_helper.from_array(np.array([1.0]), name="const_value"),
+            },
+            required_initializers={},
+        )
+
+    def get_test_specs(self) -> list:
+        return [
+            valid_test("basic")
+            .description("Basic Constant node returning scalar 1.0")
+            .tags("basic", "constant")
+            .build(),
+            e2e_test("e2e_basic")
+            .description("End-to-end test for Constant node")
+            .override_output_shapes(constant_output=[1])
+            .tags("e2e", "constant")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers/conv_config.py

@@ -0,0 +1,154 @@
+import numpy as np
+
+from python.core.model_processing.onnx_quantizer.exceptions import InvalidParamError
+from python.tests.onnx_quantizer_tests import TEST_RNG_SEED
+from python.tests.onnx_quantizer_tests.layers.base import (
+    LayerTestSpec,
+    e2e_test,
+    error_test,
+    valid_test,
+)
+from python.tests.onnx_quantizer_tests.layers.factory import (
+    BaseLayerConfigProvider,
+    LayerTestConfig,
+)
+
+
+class ConvConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Conv layers"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Conv"
+
+    def get_config(self) -> LayerTestConfig:
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        return LayerTestConfig(
+            op_type="Conv",
+            valid_inputs=["input", "conv_weight", "conv_bias"],
+            valid_attributes={
+                "strides": [1, 1],
+                "kernel_shape": [3, 3],
+                "dilations": [1, 1],
+                "pads": [1, 1, 1, 1],
+            },
+            required_initializers={
+                "conv_weight": rng.normal(0, 1, (32, 16, 3, 3)),
+                "conv_bias": rng.normal(0, 1, 32),
+            },
+        )
+
+    def get_test_specs(self) -> list[LayerTestSpec]:
+        """Return all test specifications for Conv layers"""
+        rng = np.random.default_rng(TEST_RNG_SEED)
+        return [
+            # Valid variations
+            valid_test("basic")
+            .description("Basic 2D convolution")
+            .tags("basic", "2d")
+            .build(),
+            valid_test("different_padding")
+            .description("Convolution with different padding")
+            .override_attrs(pads=[2, 2, 2, 2], kernel_shape=[5, 5])
+            .override_initializer("conv_weight", rng.normal(0, 1, (32, 16, 5, 5)))
+            .tags("padding", "5x5_kernel")
+            .build(),
+            # E2E test
+            e2e_test("e2e_basic")
+            .description("End-to-end test for basic 2D convolution")
+            .override_input_shapes(input=[1, 3, 4, 4])
+            .override_output_shapes(conv_output=[1, 8, 4, 4])
+            .override_initializer("conv_weight", rng.normal(0, 1, (8, 3, 3, 3)))
+            .override_initializer("conv_bias", rng.normal(0, 1, 8))
+            .tags("e2e", "basic", "2d")
+            .build(),
+            # Error cases
+            error_test("no_bias")
+            .description("2D convolution without bias")
+            .override_inputs("input", "conv_weight")
+            .override_attrs(strides=[2, 2], kernel_shape=[5, 5])
+            .override_initializer("conv_weight", rng.normal(0, 1, (64, 16, 5, 5)))
+            .expects_error(
+                InvalidParamError,
+                "Expected at least 3 inputs (input, weights, bias), got 2",
+            )
+            .tags("no_bias", "stride_2")
+            .build(),
+            error_test("conv3d_unsupported")
+            .description("3D convolution should raise error")
+            .override_attrs(
+                kernel_shape=[3, 3, 3],
+                strides=[1, 1, 1],
+                dilations=[1, 1, 1],
+                pads=[1, 1, 1, 1, 1, 1],
+            )
+            .override_initializer(
+                "conv_weight",
+                rng.normal(0, 1, (32, 16, 3, 3, 3)),
+            )
+            .expects_error(
+                InvalidParamError,
+                "Unsupported Conv weight dimensionality 5",
+            )
+            .tags("3d", "unsupported")
+            .build(),
+            error_test("invalid_stride")
+            .description("Invalid stride values")
+            .override_attrs(strides=[0, 1])
+            .override_inputs("input", "conv_weight")
+            .expects_error(InvalidParamError, "stride must be positive")
+            .tags("invalid_params")
+            .skip("Not yet supported")
+            .build(),
+            error_test("negative_dilation")
+            .description("Negative dilation values")
+            .override_attrs(dilations=[-1, 1])
+            .expects_error(InvalidParamError, "dilation must be positive")
+            .tags("invalid_params")
+            .skip("Not yet supported")
+            .build(),
+            error_test("invalid_kernel_shape_long")
+            .description("kernel_shape too long (length 3)")
+            .override_attrs(kernel_shape=[3, 3, 3])
+            .override_initializer(
+                "conv_weight",
+                rng.normal(0, 1, (32, 16, 3, 3, 3)),
+            )
+            .expects_error(InvalidParamError, "kernel_shape")
+            .tags("invalid_attr_length")
+            .build(),
+            # Missing required attributes
+            error_test("missing_strides")
+            .description("Conv node missing 'strides' attribute")
+            .omit_attrs("strides")  # exclude strides
+            .override_attrs(
+                kernel_shape=[3, 3],
+                dilations=[1, 1],
+                pads=[1, 1, 1, 1],
+            )  # supply others explicitly
+            .expects_error(InvalidParamError, "strides")
+            .tags("missing_attr")
+            .build(),
+            error_test("missing_kernel_shape")
+            .description("Conv node missing 'kernel_shape' attribute")
+            .omit_attrs("kernel_shape")  # exclude kernel_shape
+            .override_attrs(
+                strides=[3, 3],
+                dilations=[1, 1],
+                pads=[1, 1, 1, 1],
+            )  # supply others explicitly
+            .expects_error(InvalidParamError, "kernel_shape")
+            .tags("missing_attr")
+            .build(),
+            error_test("missing_dilations")
+            .description("Conv node missing 'dilations' attribute")
+            .omit_attrs("dilations")  # exclude dilations
+            .override_attrs(
+                strides=[3, 3],
+                kernel_shape=[3, 3],
+                pads=[1, 1, 1, 1],
+            )  # supply others explicitly
+            .expects_error(InvalidParamError, "dilations")
+            .tags("missing_attr")
+            .build(),
+        ]

python/tests/onnx_quantizer_tests/layers/factory.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import importlib
+import inspect
+import logging
+import typing
+from pathlib import Path
+
+from .base import BaseLayerConfigProvider, LayerTestConfig, LayerTestSpec, SpecType
+
+logger = logging.getLogger(__name__)
+
+
+class TestLayerFactory:
+    """Enhanced factory for creating test configurations for different layer types"""
+
+    _providers: typing.ClassVar = {}
+    _initialized = False
+
+    @classmethod
+    def _discover_providers(cls: TestLayerFactory) -> None:
+        """Automatically discover all BaseLayerConfigProvider subclasses"""
+        if cls._initialized:
+            return
+
+        current_dir = Path(__file__).parent
+        config_files = [
+            f.stem
+            for f in Path(current_dir).iterdir()
+            if f.is_file() and f.name.endswith("_config.py") and f.name != "__init__.py"
+        ]
+        logger.debug("Discovered config files: %s", config_files)
+
+        for module_name in config_files:
+            try:
+                module = importlib.import_module(f".{module_name}", package=__package__)
+
+                for _, obj in inspect.getmembers(module, inspect.isclass):
+                    if (
+                        issubclass(obj, BaseLayerConfigProvider)
+                        and obj is not BaseLayerConfigProvider
+                    ):
+
+                        provider_instance = obj()
+                        cls._providers[provider_instance.layer_name] = provider_instance
+
+            except ImportError as e:  # noqa: PERF203
+                msg = f"Warning: Could not import {module_name}: {e}"
+                logger.warning(msg)
+
+        cls._initialized = True
+
+    # Existing methods (keep your current implementation)
+    @classmethod
+    def get_layer_configs(cls) -> dict[str, LayerTestConfig]:
+        """Get test configurations for all supported layers"""
+        cls._discover_providers()
+        logger.debug("Retrieved layer configs: %s", list(cls._providers.keys()))
+        return {
+            name: provider.get_config() for name, provider in cls._providers.items()
+        }
+
+    @classmethod
+    def get_layer_config(cls, layer_name: str) -> LayerTestConfig:
+        """Get test configuration for a specific layer"""
+        cls._discover_providers()
+        if layer_name not in cls._providers:
+            msg = f"No config provider found for layer: {layer_name}"
+            raise ValueError(msg)
+        return cls._providers[layer_name].get_config()
+
+    @classmethod
+    def get_available_layers(cls) -> list[str]:
+        """Get list of all available layer types"""
+        cls._discover_providers()
+        return list(cls._providers.keys())
+
+    @classmethod
+    def register_provider(cls, provider: BaseLayerConfigProvider) -> None:
+        """Register a new config provider"""
+        cls._providers[provider.layer_name] = provider
+
+    # NEW enhanced methods for test specifications
+    @classmethod
+    def get_all_test_cases(cls) -> list[tuple[str, LayerTestConfig, LayerTestSpec]]:
+        """Get all test cases as (layer_name, config, test_spec) tuples"""
+        cls._discover_providers()
+        test_cases = []
+
+        for layer_name, provider in cls._providers.items():
+            config = provider.get_config()
+            test_specs = provider.get_test_specs()
+
+            # If no test specs defined, create a basic valid test
+            if not test_specs:
+
+                test_specs = [LayerTestSpec("basic", SpecType.VALID, "Basic test")]
+
+            test_cases.extend((layer_name, config, spec) for spec in test_specs)
+
+        return test_cases
+
+    @classmethod
+    def get_test_cases_by_type(
+        cls,
+        test_type: SpecType,
+    ) -> list[tuple[str, LayerTestConfig, LayerTestSpec]]:
+        """Get test cases of a specific type"""
+        return [
+            (layer, config, spec)
+            for layer, config, spec in cls.get_all_test_cases()
+            if spec.spec_type == test_type
+        ]
+
+    @classmethod
+    def get_test_cases_by_layer(
+        cls,
+        layer_name: str,
+    ) -> list[tuple[str, LayerTestConfig, LayerTestSpec]]:
+        """Get test cases for a specific layer"""
+        return [
+            (layer, config, spec)
+            for layer, config, spec in cls.get_all_test_cases()
+            if layer == layer_name
+        ]
+
+    @classmethod
+    def get_test_cases_by_tag(
+        cls,
+        tag: str,
+    ) -> list[tuple[str, LayerTestConfig, LayerTestSpec]]:
+        """Get test cases with a specific tag"""
+        result = [
+            (layer, config, spec)
+            for layer, config, spec in cls.get_all_test_cases()
+            if tag in spec.tags
+        ]
+        logger.debug("Found tests %s", result)
+        if not result:
+            msg = f"No test cases found for tag: {tag}"
+            raise ValueError(msg)
+        return result

python/tests/onnx_quantizer_tests/layers/flatten_config.py

@@ -0,0 +1,61 @@
+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 FlattenConfigProvider(BaseLayerConfigProvider):
+    """Test configuration provider for Flatten layers"""
+
+    @property
+    def layer_name(self) -> str:
+        return "Flatten"
+
+    def get_config(self) -> LayerTestConfig:
+        return LayerTestConfig(
+            op_type="Flatten",
+            valid_inputs=["input"],
+            valid_attributes={"axis": 1},
+            required_initializers={},
+        )
+
+    def get_test_specs(self) -> list:
+        return [
+            # --- VALID TESTS ---
+            valid_test("basic")
+            .description("Basic Flatten from (1,3,4,4) to (1,48)")
+            .tags("basic", "flatten")
+            .build(),
+            valid_test("flatten_axis0")
+            .description("Flatten with axis=0 (entire tensor flattened)")
+            .override_attrs(axis=0)
+            .tags("flatten", "axis0")
+            .build(),
+            valid_test("flatten_axis2")
+            .description("Flatten starting at axis=2")
+            .override_attrs(axis=2)
+            .tags("flatten", "axis2")
+            .build(),
+            valid_test("flatten_axis3")
+            .description("Flatten starting at axis=3 (minimal flatten)")
+            .override_attrs(axis=3)
+            .tags("flatten", "axis3")
+            .build(),
+            e2e_test("e2e_basic")
+            .description("End-to-end test for Flatten layer")
+            .override_input_shapes(input=[1, 3, 4, 4])
+            .override_output_shapes(flatten_output=[1, 48])
+            .tags("e2e", "flatten")
+            .build(),
+            # --- EDGE CASE / SKIPPED TEST ---
+            valid_test("large_input")
+            .description("Large input flatten (performance test)")
+            .override_input_shapes(input=[1, 3, 256, 256])
+            .tags("flatten", "large", "performance")
+            .skip("Performance test, skipped by default")
+            .build(),
+        ]