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

python/core/model_processing/onnx_quantizer/layers/maxpool.py

@@ -1,20 +1,28 @@
 from __future__ import annotations
 
-import onnx
-from onnx import helper
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import onnx
 
 from python.core.model_processing.onnx_custom_ops.onnx_helpers import (
-    extract_attributes,
     get_attribute_ints,
 )
 from python.core.model_processing.onnx_quantizer.exceptions import InvalidParamError
 from python.core.model_processing.onnx_quantizer.layers.base import (
     BaseOpQuantizer,
+    QuantizerBase,
     ScaleConfig,
 )
 
 
-class MaxpoolQuantizer(BaseOpQuantizer):
+class QuantizeMaxpool(QuantizerBase):
+    OP_TYPE = "Int64MaxPool"
+    USE_WB = False
+    USE_SCALING = False
+
+
+class MaxpoolQuantizer(BaseOpQuantizer, QuantizeMaxpool):
     """
     Quantizer for ONNX MaxPool layers.
 
@@ -25,55 +33,26 @@ class MaxpoolQuantizer(BaseOpQuantizer):
 
     def __init__(
         self: MaxpoolQuantizer,
-        new_initializer: dict[str, onnx.TensorProto] | None = None,
+        new_initializer: list[onnx.TensorProto] | None = None,
     ) -> None:
         super().__init__()
         self.accepted_kernel_shapes = [2]
         _ = new_initializer
 
     def quantize(
-        self: BaseOpQuantizer,
+        self: MaxpoolQuantizer,
         node: onnx.NodeProto,
         graph: onnx.GraphProto,
         scale_config: ScaleConfig,
         initializer_map: dict[str, onnx.TensorProto],
     ) -> list[onnx.NodeProto]:
-        """
-        Quantize a node by converting the node to Int64 version
-
-        Args:
-            node (onnx.NodeProto): The node to quantize.
-            rescale (bool): Whether rescaling is enabled
-                (Doesnt have an affect on this op type)
-            graph (onnx.GraphProto): The ONNX graph.
-            scale_exponent (int): Scale exponent.
-            scale_base (int): The base of scaling.
-            initializer_map (dict[str, onnx.TensorProto]):
-                Map of initializer names to tensor data.
-
-        Returns:
-            List[onnx.NodeProto]: A list of ONNX nodes
-                (quantized MaxPool and any auxiliary nodes).
-        """
-        _ = initializer_map, graph
-
-        attrs = extract_attributes(node)
-        attrs["rescale"] = int(scale_config.rescale)
-
-        attr_str = {
-            k: ",".join(map(str, v)) if isinstance(v, list) else str(v)
-            for k, v in attrs.items()
-        }
-        return [
-            helper.make_node(
-                "Int64MaxPool",
-                inputs=node.input,
-                outputs=node.output,
-                name=node.name,
-                domain="ai.onnx.contrib",
-                **attr_str,
-            ),
-        ]
+        return QuantizeMaxpool.quantize(
+            self,
+            node,
+            graph,
+            scale_config,
+            initializer_map,
+        )
 
     def check_supported(
         self: MaxpoolQuantizer,
@@ -95,6 +74,7 @@ class MaxpoolQuantizer(BaseOpQuantizer):
         _ = initializer_map
         self.check_all_params_exist(node)
         self.check_params_size(node)
+        self.check_pool_pads(node)
 
     def check_all_params_exist(self: MaxpoolQuantizer, node: onnx.NodeProto) -> None:
         """Checks all parameters that are needed, do exist
@@ -131,10 +111,40 @@ class MaxpoolQuantizer(BaseOpQuantizer):
             InvalidParamError: If shape requirement is not met.
         """
 
-        kernel_shape = get_attribute_ints(node, "kernel_shape", default="N/A")
+        kernel_shape = get_attribute_ints(node, "kernel_shape", default=[])
         if len(kernel_shape) not in self.accepted_kernel_shapes:
             raise InvalidParamError(
                 node.name,
                 node.op_type,
                 f"Currently only maxpool2d is supported. Found {len(kernel_shape)}D",
             )
+
+    def check_pool_pads(self: MaxpoolQuantizer, node: onnx.NodeProto) -> None:
+        kernel_shape = get_attribute_ints(node, "kernel_shape", default=[])
+        pads = get_attribute_ints(node, "pads", default=None)
+        if pads is None:
+            return
+        num_dims = len(kernel_shape)
+        if len(pads) != num_dims * 2:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"Expected {num_dims * 2} pads, got {len(pads)}",
+            )
+
+        for dim in range(num_dims):
+            pad_before = pads[dim]
+            pad_after = pads[dim + num_dims]
+            kernel = kernel_shape[dim]
+            if pad_before >= kernel:
+                raise InvalidParamError(
+                    node.name,
+                    node.op_type,
+                    f"pads[{dim}]={pad_before} >= kernel[{dim}]={kernel}",
+                )
+            if pad_after >= kernel:
+                raise InvalidParamError(
+                    node.name,
+                    node.op_type,
+                    f"pads[{dim + num_dims}]={pad_after} >= kernel[{dim}]={kernel}",
+                )

python/core/model_processing/onnx_quantizer/layers/mul.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
+
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeMul(QuantizerBase):
+    OP_TYPE = "Int64Mul"
+    USE_WB = True
+    USE_SCALING = True
+    SCALE_PLAN: ClassVar = {0: 1, 1: 1}
+
+
+class MulQuantizer(BaseOpQuantizer, QuantizeMul):
+    """
+    Quantizer for ONNX Mul layers.
+
+    - Uses custom Mul layer to incorporate rescaling, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: MulQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: MulQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        return QuantizeMul.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: MulQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        pass

python/core/model_processing/onnx_quantizer/layers/relu.py

@@ -1,14 +1,24 @@
 from __future__ import annotations
 
-import onnx
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from onnx import GraphProto, NodeProto, TensorProto
 
 from python.core.model_processing.onnx_quantizer.layers.base import (
     BaseOpQuantizer,
+    QuantizerBase,
     ScaleConfig,
 )
 
 
-class ReluQuantizer(BaseOpQuantizer):
+class QuantizeRelu(QuantizerBase):
+    OP_TYPE = "Int64Relu"
+    USE_WB = False
+    USE_SCALING = False
+
+
+class ReluQuantizer(BaseOpQuantizer, QuantizeRelu):
     """
     Quantizer for ONNX ReLU layers.
 
@@ -19,49 +29,24 @@ class ReluQuantizer(BaseOpQuantizer):
 
     def __init__(
         self: ReluQuantizer,
-        new_initializer: dict[str, onnx.TensorProto] | None = None,
+        new_initializer: list[TensorProto] | None = None,
     ) -> None:
         super().__init__()
         _ = new_initializer
 
     def quantize(
         self: ReluQuantizer,
-        node: onnx.NodeProto,
-        graph: onnx.GraphProto,
+        node: NodeProto,
+        graph: GraphProto,
         scale_config: ScaleConfig,
-        initializer_map: dict[str, onnx.TensorProto],
-    ) -> list[onnx.NodeProto]:
-        """
-        Quantize a node by converting the node to Int64 version
-
-        Args:
-            node (onnx.NodeProto): The node to quantize.
-            rescale (bool): Whether rescaling is enabled
-                (Doesnt have an affect on this op type)
-            graph (onnx.GraphProto): The ONNX graph.
-            scale_exponent (int): Scale exponent.
-            scale_base (int): The base of scaling.
-            initializer_map (dict[str, onnx.TensorProto]):
-                Map of initializer names to tensor data.
-
-        Returns:
-            List[onnx.NodeProto]: The quantized ONNX node.
-        """
-        _ = graph, scale_config, initializer_map
-        return [
-            onnx.helper.make_node(
-                "Int64Relu",
-                inputs=node.input,
-                outputs=node.output,  # preserve original output name
-                name=node.name,
-                domain="ai.onnx.contrib",
-            ),
-        ]
+        initializer_map: dict[str, TensorProto],
+    ) -> list[NodeProto]:
+        return QuantizeRelu.quantize(self, node, graph, scale_config, initializer_map)
 
     def check_supported(
         self: ReluQuantizer,
-        node: onnx.NodeProto,
-        initializer_map: dict[str, onnx.TensorProto] | None = None,
+        node: NodeProto,
+        initializer_map: dict[str, TensorProto] | None = None,
     ) -> None:
         """
         Perform high-level validation to ensure that this node

python/core/model_processing/onnx_quantizer/layers/sub.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
+
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeSub(QuantizerBase):
+    OP_TYPE = "Sub"
+    DOMAIN = ""
+    USE_WB = True
+    USE_SCALING = False
+    SCALE_PLAN: ClassVar = {0: 1, 1: 1}
+
+
+class SubQuantizer(BaseOpQuantizer, QuantizeSub):
+    """
+    Quantizer for ONNX Sub layers.
+
+    - Uses standard ONNX Sub layer in standard domain, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: SubQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: SubQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        return QuantizeSub.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: SubQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        pass

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
-from typing import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
 import onnx
 
@@ -9,17 +12,23 @@ from python.core.model_processing.onnx_quantizer.exceptions import (
     MissingHandlerError,
     UnsupportedOpError,
 )
+from python.core.model_processing.onnx_quantizer.layers.add import AddQuantizer
 from python.core.model_processing.onnx_quantizer.layers.base import (
     PassthroughQuantizer,
     ScaleConfig,
 )
+from python.core.model_processing.onnx_quantizer.layers.batchnorm import (
+    BatchnormQuantizer,
+)
 from python.core.model_processing.onnx_quantizer.layers.constant import (
     ConstantQuantizer,
 )
 from python.core.model_processing.onnx_quantizer.layers.conv import ConvQuantizer
 from python.core.model_processing.onnx_quantizer.layers.gemm import GemmQuantizer
 from python.core.model_processing.onnx_quantizer.layers.maxpool import MaxpoolQuantizer
+from python.core.model_processing.onnx_quantizer.layers.mul import MulQuantizer
 from python.core.model_processing.onnx_quantizer.layers.relu import ReluQuantizer
+from python.core.model_processing.onnx_quantizer.layers.sub import SubQuantizer
 
 
 class ONNXOpQuantizer:
@@ -64,6 +73,9 @@ class ONNXOpQuantizer:
         self.new_initializers = []
 
         # Register handlers
+        self.register("Add", AddQuantizer(self.new_initializers))
+        self.register("Sub", SubQuantizer(self.new_initializers))
+        self.register("Mul", MulQuantizer(self.new_initializers))
         self.register("Conv", ConvQuantizer(self.new_initializers))
         self.register("Relu", ReluQuantizer())
         self.register("Reshape", PassthroughQuantizer())
@@ -71,6 +83,7 @@ class ONNXOpQuantizer:
         self.register("Constant", ConstantQuantizer())
         self.register("MaxPool", MaxpoolQuantizer())
         self.register("Flatten", PassthroughQuantizer())
+        self.register("BatchNormalization", BatchnormQuantizer(self.new_initializers))
 
     def register(
         self: ONNXOpQuantizer,
@@ -198,3 +211,32 @@ class ONNXOpQuantizer:
             dict[str, onnx.TensorProto]: Map from initializer name to tensors in graph.
         """
         return {init.name: init for init in model.graph.initializer}
+
+    def apply_pre_analysis_transforms(
+        self: ONNXOpQuantizer,
+        model: onnx.ModelProto,
+        scale_exponent: int,
+        scale_base: int,
+    ) -> onnx.ModelProto:
+        """
+        Give each registered handler a chance to rewrite the model before analysis.
+        """
+        graph = model.graph
+        initializer_map = self.get_initializer_map(model)
+
+        # We allow handlers to modify graph in-place.
+        # (Nodes may be replaced, removed, or new nodes added.)
+        for node in list(graph.node):
+            handler = self.handlers.get(node.op_type)
+            if handler and hasattr(handler, "pre_analysis_transform"):
+                handler.pre_analysis_transform(
+                    node,
+                    graph,
+                    initializer_map,
+                    scale_exponent=scale_exponent,
+                    scale_base=scale_base,
+                )
+            # Refresh map if transforms may add initializers
+            initializer_map = self.get_initializer_map(model)
+
+        return model

python/core/utils/general_layer_functions.py

@@ -83,7 +83,9 @@ class GeneralLayerFunctions:
             tensor = torch.mul(tensor, self.scale_base**self.scale_exponent)
 
         tensor = tensor.long()
+        return self.reshape_inputs(tensor)
 
+    def reshape_inputs(self: GeneralLayerFunctions, tensor: torch.Tensor) -> None:
         if hasattr(self, "input_shape"):
             shape = self.input_shape
             if hasattr(self, "adjust_shape") and callable(
@@ -94,6 +96,7 @@ class GeneralLayerFunctions:
                 tensor = tensor.reshape(shape)
             except RuntimeError as e:
                 raise ShapeMismatchError(shape, list(tensor.shape)) from e
+
         return tensor
 
     def get_inputs(
@@ -154,21 +157,19 @@ class GeneralLayerFunctions:
                 # If unknown dim in batch spot, assume batch size of 1
                 first_key = next(iter(keys))
                 input_shape = self.input_shape[first_key]
-                input_shape[0] = 1 if input_shape[0] < 1 else input_shape[0]
+                input_shape[0] = max(input_shape[0], 1)
                 return self.get_rand_inputs(input_shape)
             inputs = {}
             for key in keys:
                 # If unknown dim in batch spot, assume batch size of 1
-                input_shape = self.input_shape[keys[key]]
-                if not isinstance(input_shape, list) and not isinstance(
-                    input_shape,
-                    tuple,
-                ):
+                input_shape = self.input_shape[key]
+                if not isinstance(input_shape, (list, tuple)):
                     msg = f"Invalid input shape for key '{key}': {input_shape}"
                     raise CircuitUtilsError(msg)
-                input_shape[0] = 1 if input_shape[0] < 1 else input_shape[0]
+                input_shape[0] = max(input_shape[0], 1)
                 inputs[key] = self.get_rand_inputs(input_shape)
             return inputs
+
         if not (hasattr(self, "scale_base") and hasattr(self, "scale_exponent")):
             attr_name = "scale_base/scale_exponent"
             context = "needed for scaling random inputs"
@@ -195,8 +196,10 @@ class GeneralLayerFunctions:
             torch.Tensor: A tensor of random values in [-1, 1).
         """
         if not isinstance(input_shape, (list, tuple)):
-            msg = f"Invalid input_shape type: {type(input_shape)}."
-            " Expected list or tuple of ints."
+            msg = (
+                f"Invalid input_shape type: {type(input_shape)}."
+                " Expected list or tuple of ints."
+            )
             raise CircuitUtilsError(msg)
         if not all(isinstance(x, int) and x > 0 for x in input_shape):
             raise ShapeMismatchError(
@@ -240,9 +243,11 @@ class GeneralLayerFunctions:
             try:
                 rescaled = torch.div(outputs, self.scale_base**self.scale_exponent)
             except Exception as e:
-                msg = "Failed to rescale outputs using scale_base="
-                f"{getattr(self, 'scale_base', None)} "
-                f"and scale_exponent={getattr(self, 'scale_exponent', None)}: {e}"
+                msg = (
+                    "Failed to rescale outputs using scale_base="
+                    f"{getattr(self, 'scale_base', None)} "
+                    f"and scale_exponent={getattr(self, 'scale_exponent', None)}: {e}"
+                )
                 raise CircuitUtilsError(msg) from e
             return {
                 "output": outputs.long().tolist(),

python/core/utils/model_registry.py

@@ -147,7 +147,8 @@ def get_models_to_test(
 
     Args:
         selected_models (list[str], optional):
-            A list of model names to include. Defaults to None.
+            A list of model names to include. If None, returns empty list.
+            If provided but no matches found, returns empty list. Defaults to None.
         source_filter (str, optional):
             Restrict models to a specific source (e.g., "onnx", "class").
             Defaults to None.
@@ -155,10 +156,12 @@ def get_models_to_test(
     Returns:
         list[ModelEntry]: A filtered list of model entries.
     """
+    if selected_models is None:
+        return []
+
     models = MODELS_TO_TEST
 
-    if selected_models is not None:
-        models = [m for m in models if m.name in selected_models]
+    models = [m for m in models if m.name in selected_models]
 
     if source_filter is not None:
         models = [m for m in models if m.source == source_filter]

python/scripts/gen_and_bench.py

@@ -247,12 +247,12 @@ def export_onnx(
 
 
 def write_input_json(json_path: Path, input_shape: tuple[int] = (1, 4, 28, 28)) -> None:
-    """Write a zero-valued input tensor to JSON alongside its [N,C,H,W] shape."""
+    """Write a zero-valued input tensor to JSON without shape information."""
     json_path.parent.mkdir(parents=True, exist_ok=True)
     n, c, h, w = input_shape
     arr = [0.0] * (n * c * h * w)
     with json_path.open("w", encoding="utf-8") as f:
-        json.dump({"input": arr, "shape": [n, c, h, w]}, f)
+        json.dump({"input": arr}, f)
 
 
 def run_bench(