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

python/core/model_processing/onnx_quantizer/layers/add.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 QuantizeAdd(QuantizerBase):
+    OP_TYPE = "Add"
+    DOMAIN = ""
+    USE_WB = True
+    USE_SCALING = False
+    SCALE_PLAN: ClassVar = {0: 1, 1: 1}
+
+
+class AddQuantizer(BaseOpQuantizer, QuantizeAdd):
+    """
+    Quantizer for ONNX Add layers.
+
+    - Uses standard ONNX Add layer in standard domain, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: AddQuantizer,
+        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: AddQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        return QuantizeAdd.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: AddQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        pass

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

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import ClassVar
 
 import numpy as np
 import onnx
 from onnx import helper, numpy_helper
 
 from python.core.model_processing.onnx_custom_ops.onnx_helpers import (
+    extract_attributes,
     replace_input_references,
 )
 from python.core.model_processing.onnx_quantizer.exceptions import (
@@ -188,7 +190,7 @@ class BaseOpQuantizer:
         """
         self.validate_node_has_output(node)
 
-        original_output = node.output.get(0)
+        original_output = node.output[0]
         quantized_output = original_output + "_raw"
         node.output[0] = quantized_output
 
@@ -294,6 +296,61 @@ class BaseOpQuantizer:
         # === Mutate the original node ===
         return nodes, new_inputs
 
+    def add_scaled_initializer_inputs(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto],
+        scale_base: int,
+        scale_exponent: int,
+        scale_plan: dict[int, int],
+    ) -> tuple[list[onnx.NodeProto], list[str]]:
+        """
+        Scale and cast specific initializer inputs
+        of a node according to a scaling plan.
+
+        Handles optional inputs gracefully (e.g. missing bias).
+        """
+        new_nodes: list[onnx.NodeProto] = []
+        new_inputs = list(node.input)
+
+        for input_idx, scale_mult in scale_plan.items():
+            # Skip if node doesn't have that many inputs (e.g. missing bias)
+            if input_idx >= len(node.input):
+                # Just ignore — optional input not provided
+                continue
+
+            input_name = node.input[input_idx]
+            if not input_name:
+                # Empty input name → optional input not present
+                continue
+
+            if input_name not in initializer_map:
+                # Optional inputs may be missing from initializers (e.g., dynamic bias)
+                continue
+
+            tensor = initializer_map[input_name]
+            if not tensor.name:
+                raise HandlerImplementationError(
+                    op_type=node.op_type,
+                    message=f"Initializer tensor for '{input_name}' on node "
+                    f"'{node.name}' is missing a name.",
+                )
+
+            # Scale according to plan (e.g., scale_exponent * 2 for bias)
+            quant_name, mul_node, cast_node = self.insert_scale_node(
+                tensor=tensor,
+                scale_base=scale_base,
+                scale_exponent=(scale_exponent * scale_mult),
+            )
+
+            # Update node input to point to scaled version
+            new_inputs[input_idx] = quant_name
+
+            # Record new scaling/cast nodes
+            new_nodes.extend([mul_node, cast_node])
+
+        return new_nodes, new_inputs
+
     def insert_scale_node(
         self: BaseOpQuantizer,
         tensor: onnx.TensorProto,
@@ -360,6 +417,69 @@ class BaseOpQuantizer:
         return output_name, mul_node, cast_to_int64
 
 
+class QuantizerBase:
+    OP_TYPE = None
+    DOMAIN = "ai.onnx.contrib"
+    DEFAULT_ATTRS: ClassVar = {}
+    USE_WB = False
+    USE_SCALING = False
+
+    def quantize(
+        self,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        """Generic quantization template for most Int64 ops."""
+        _ = graph
+        nodes = []
+
+        # (1) Quantize weights/bias if applicable
+        if self.USE_WB:
+            # Each subclass defines its scaling plan for which inputs get scaled and how
+            scale_plan = getattr(self, "SCALE_PLAN", {1: 1, 2: 2})  # default for W & B
+            nodes, new_inputs = self.add_scaled_initializer_inputs(
+                node=node,
+                initializer_map=initializer_map,
+                scale_base=scale_config.base,
+                scale_exponent=scale_config.exponent,
+                scale_plan=scale_plan,
+            )
+            node.input[:] = new_inputs
+
+        # (2) Collect & merge attributes
+        attrs = extract_attributes(node)
+        for k, v in self.DEFAULT_ATTRS.items():
+            attrs.setdefault(k, v)
+        if self.USE_SCALING:
+            attrs["rescale"] = int(scale_config.rescale)
+
+        # (3) Add scaling constant if needed
+        if self.USE_SCALING:
+            scale_value = self.get_scaling(scale_config.base, scale_config.exponent)
+            scale_name = f"{node.name}_int_scaler"
+            scale_tensor = numpy_helper.from_array(
+                np.array([scale_value], dtype=np.int64),
+                name=scale_name,
+            )
+            self.new_initializers.append(scale_tensor)
+            node.input.append(scale_name)
+
+        # (4) Create quantized node
+        quantized_node = onnx.helper.make_node(
+            self.OP_TYPE,
+            inputs=node.input,
+            outputs=node.output,
+            name=node.name,
+            domain=self.DOMAIN,
+            **attrs,
+        )
+
+        nodes.append(quantized_node)
+        return nodes
+
+
 class PassthroughQuantizer(BaseOpQuantizer):
     """
     Quantizer that leaves the node unchanged.

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

@@ -38,7 +38,7 @@ class ConstantQuantizer(BaseOpQuantizer):
 
     def __init__(
         self: ConstantQuantizer,
-        new_initializer: dict[str, onnx.TensorProto] | None = None,
+        new_initializer: list[onnx.TensorProto] | None = None,
     ) -> None:
         super().__init__()
         _ = new_initializer

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

@@ -1,18 +1,27 @@
 from __future__ import annotations
 
-import numpy as np
-import onnx
-from onnx import numpy_helper
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
 
-from python.core.model_processing.onnx_custom_ops.onnx_helpers import extract_attributes
 from python.core.model_processing.onnx_quantizer.exceptions import InvalidParamError
 from python.core.model_processing.onnx_quantizer.layers.base import (
     BaseOpQuantizer,
+    QuantizerBase,
     ScaleConfig,
 )
 
 
-class ConvQuantizer(BaseOpQuantizer):
+class QuantizeConv(QuantizerBase):
+    OP_TYPE = "Int64Conv"
+    USE_WB = True
+    USE_SCALING = True
+    DEFAULT_ATTRS: ClassVar = {"group": 1, "auto_pad": "NOTSET"}
+    SCALE_PLAN: ClassVar = {1: 1, 2: 2}  # weight = 1x scale, bias = 2x scale
+
+
+class ConvQuantizer(BaseOpQuantizer, QuantizeConv):
     """
     Quantizer for ONNX Conv layers.
 
@@ -23,9 +32,12 @@ class ConvQuantizer(BaseOpQuantizer):
 
     def __init__(
         self: ConvQuantizer,
-        new_initializers: dict[str, onnx.TensorProto],
+        new_initializers: list[onnx.TensorProto] | None = None,
     ) -> None:
-        self.new_initializers = new_initializers
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
 
     def quantize(
         self: ConvQuantizer,
@@ -34,67 +46,7 @@ class ConvQuantizer(BaseOpQuantizer):
         scale_config: ScaleConfig,
         initializer_map: dict[str, onnx.TensorProto],
     ) -> list[onnx.NodeProto]:
-        """
-        Quantize a Conv node by:
-        1. Quantizing its weights and bias.
-        2. Adding a scale constant.
-        3. Replacing it with an Int64Conv node.
-
-        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 and any auxiliary nodes).
-        """
-        _ = graph
-
-        nodes = []
-        output_name = f"{node.name}_int"
-
-        nodes, node.input[:] = self.add_nodes_w_and_b(
-            node=node,
-            scale_exponent=scale_config.exponent,
-            scale_base=scale_config.base,
-            initializer_map=initializer_map,
-        )
-        attrs = extract_attributes(node)
-        attrs.setdefault("group", 1)
-        attrs.setdefault("auto_pad", "NOTSET")
-
-        attrs["rescale"] = int(scale_config.rescale)
-
-        scale_value = self.get_scaling(
-            scale_config.base,
-            scale_config.exponent,
-        )
-
-        # Create scale constant
-        scale_const_name = f"{output_name}_scaler"
-        scale_tensor = numpy_helper.from_array(
-            np.array([scale_value], dtype=np.int64),
-            name=scale_const_name,
-        )
-        self.new_initializers.append(scale_tensor)
-        node.input.append(scale_const_name)
-        int64_conv_node = onnx.helper.make_node(
-            "Int64Conv",
-            inputs=node.input,
-            outputs=node.output,  # preserve original output name
-            name=node.name,
-            domain="ai.onnx.contrib",
-            **attrs,
-        )
-
-        nodes.append(int64_conv_node)
-        return nodes
+        return QuantizeConv.quantize(self, node, graph, scale_config, initializer_map)
 
     def check_supported(
         self: ConvQuantizer,

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

@@ -1,18 +1,27 @@
 from __future__ import annotations
 
-import numpy as np
-import onnx
-from onnx import numpy_helper
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
 
-from python.core.model_processing.onnx_custom_ops.onnx_helpers import extract_attributes
 from python.core.model_processing.onnx_quantizer.exceptions import InvalidParamError
 from python.core.model_processing.onnx_quantizer.layers.base import (
     BaseOpQuantizer,
+    QuantizerBase,
     ScaleConfig,
 )
 
 
-class GemmQuantizer(BaseOpQuantizer):
+class QuantizeGemm(QuantizerBase):
+    OP_TYPE = "Int64Gemm"
+    USE_WB = True
+    USE_SCALING = True
+    DEFAULT_ATTRS: ClassVar = {"transA": 0, "transB": 0}
+    SCALE_PLAN: ClassVar = {1: 1, 2: 2}
+
+
+class GemmQuantizer(BaseOpQuantizer, QuantizeGemm):
     """
     Quantizer for ONNX Gemm layers.
 
@@ -23,9 +32,12 @@ class GemmQuantizer(BaseOpQuantizer):
 
     def __init__(
         self: GemmQuantizer,
-        new_initializers: dict[str, onnx.TensorProto],
+        new_initializers: list[onnx.TensorProto] | None = None,
     ) -> None:
-        self.new_initializers = new_initializers
+        super().__init__()
+        # Only replace if caller provided something
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
 
     def quantize(
         self: GemmQuantizer,
@@ -34,65 +46,7 @@ class GemmQuantizer(BaseOpQuantizer):
         scale_config: ScaleConfig,
         initializer_map: dict[str, onnx.TensorProto],
     ) -> list[onnx.NodeProto]:
-        """
-        Quantize a Gemm node by:
-        1. Quantizing its weights and bias.
-        2. Adding a scale constant.
-        3. Replacing it with an Int64Gemm node.
-
-        Args:
-            node (onnx.NodeProto): The node to quantize.
-            rescale (bool): Whether rescaling is enabled
-            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 and any auxiliary nodes).
-        """
-        _ = graph
-        nodes = []
-        output_name = f"{node.name}_int"
-
-        nodes, new_inputs = self.add_nodes_w_and_b(
-            node=node,
-            scale_exponent=scale_config.exponent,
-            scale_base=scale_config.base,
-            initializer_map=initializer_map,
-        )
-        node.input[:] = new_inputs
-
-        attrs = extract_attributes(node)
-        attrs.setdefault("transA", 0)
-        attrs.setdefault("transB", 0)
-        attrs["rescale"] = int(scale_config.rescale)
-
-        scale_value = self.get_scaling(
-            scale_config.base,
-            scale_config.exponent,
-        )
-
-        # === Create scale constant ===
-        scale_const_name = f"{output_name}_scaler"
-        scale_tensor = numpy_helper.from_array(
-            np.array([scale_value], dtype=np.int64),
-            name=scale_const_name,
-        )
-        self.new_initializers.append(scale_tensor)
-        node.input.append(scale_const_name)
-        int64_gemm = onnx.helper.make_node(
-            "Int64Gemm",
-            inputs=node.input,
-            outputs=node.output,  # preserve original output name
-            name=output_name,
-            domain="ai.onnx.contrib",
-            **attrs,
-        )
-        nodes.append(int64_gemm)
-        return nodes
+        return QuantizeGemm.quantize(self, node, graph, scale_config, initializer_map)
 
     def check_supported(
         self: GemmQuantizer,

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/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