JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl

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

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+import onnx
+from onnx import numpy_helper
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    HandlerImplementationError,
+)
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    ScaleConfig,
+)
+
+
+class ConstantQuantizer(BaseOpQuantizer):
+    """
+    Quantizer for ONNX Constant node.
+
+    This quantizer only modifies constants that are:
+    - Numeric tensors
+    - Used directly in computation
+
+    Constants used for shape, indexing, or other non-numeric roles are left unchanged.
+    """
+
+    DATA_OPS: ClassVar = {
+        "Add",
+        "Mul",
+        "Conv",
+        "MatMul",
+        "Sub",
+        "Div",
+        "Gemm",
+    }  # ops that consume numeric constants
+
+    def __init__(
+        self: ConstantQuantizer,
+        new_initializer: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        _ = new_initializer
+
+    def quantize(
+        self: ConstantQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        """Apply quantization scaling to a constant if it is used in
+            numeric computation.
+
+        Args:
+            node (onnx.NodeProto): The Constant node to quantize.
+            rescale (bool): Whether rescaling is enabled
+                (Doesnt have an affect on this op type in some cases)
+            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 modified node (possibly unchanged).
+
+        Raises:
+            HandlerImplementationError: If tensor is unreadable
+        """
+        _ = initializer_map
+        self.validate_node_has_output(node)
+
+        output_name = node.output[0]
+
+        is_data_constant = any(
+            output_name in n.input and n.op_type in self.DATA_OPS for n in graph.node
+        )
+
+        if not is_data_constant:
+            # Skip quantization for non-numeric constants
+            return [node]
+
+        # Safe to quantize: numeric constant used in computation
+        for attr in node.attribute:
+            if attr.name == "value" and attr.type == onnx.AttributeProto.TENSOR:
+                try:
+                    arr = numpy_helper.to_array(attr.t).astype(np.float64)
+                except (ValueError, Exception) as e:
+                    raise HandlerImplementationError(
+                        op_type="Constant",
+                        message="Failed to read tensor from Constant node"
+                        f" '{node.name}': {e}",
+                    ) from e
+
+                arr *= self.get_scaling(
+                    scale_config.base,
+                    scale_config.exponent,
+                )
+                attr.t.CopyFrom(numpy_helper.from_array(arr, name=""))
+
+        node.name += "_quant"
+        return [node]
+
+    def check_supported(
+        self: ConstantQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        """All Constant nodes are supported... For now.
+
+        Args:
+            node (onnx.NodeProto): Node to be checked
+            initializer_map (dict[str, onnx.TensorProto], optional):
+                Map of initializer names to tensor data. Defaults to None.
+        """
+        _ = node, initializer_map

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

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import numpy as np
+import onnx
+from onnx import numpy_helper
+
+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,
+    ScaleConfig,
+)
+
+
+class ConvQuantizer(BaseOpQuantizer):
+    """
+    Quantizer for ONNX Conv layers.
+
+    - Replaces standard Conv with Int64Conv from the `ai.onnx.contrib` domain
+        and makes relevant additional changes to the graph.
+    - Validates that all required Conv parameters are present.
+    """
+
+    def __init__(
+        self: ConvQuantizer,
+        new_initializers: dict[str, onnx.TensorProto],
+    ) -> None:
+        self.new_initializers = new_initializers
+
+    def quantize(
+        self: ConvQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        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
+
+    def check_supported(
+        self: ConvQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> None:
+        """
+        Perform high-level validation to ensure that this Conv node
+        can be quantized safely.
+
+        Args:
+            node (onnx.NodeProto): ONNX node to be checked
+            initializer_map (dict[str, onnx.TensorProto]):
+                Initializer map (name of weight or bias and tensor)
+
+        Raises:
+            InvalidParamError: If any requirement is not met.
+        """
+        num_inputs = 2
+        if len(node.input) < num_inputs:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"Expected at least 2 inputs (input, weights), got {len(node.input)}",
+            )
+        num_inputs = 3
+
+        if len(node.input) < num_inputs:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                "Expected at least 3 inputs (input, weights, bias),"
+                f" got {len(node.input)}",
+            )
+
+        self.check_supported_shape(node, initializer_map)
+        self.check_all_params_exist(node)
+
+    def check_all_params_exist(self: ConvQuantizer, node: onnx.NodeProto) -> None:
+        """Verify that all required Conv attributes are present.
+
+        Args:
+            node (onnx.NodeProto): The Conv node being validated.
+
+        Raises:
+            InvalidParamError: If any required parameter is missing.
+        """
+        # May need: ["strides", "kernel_shape", "dilations", "pads"]
+        required_attrs = ["strides", "kernel_shape", "dilations"]
+        self.validate_required_attrs(node, required_attrs)
+
+    def check_supported_shape(
+        self: ConvQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> None:
+        """Ensure that Conv weights are available and have the correct dimensionality.
+
+        Args:
+            node (onnx.NodeProto): The node being validated.
+            initializer_map (dict[str, onnx.TensorProto]):
+                Mapping of initializer tensor names to TensorProtos.
+
+        Raises:
+            InvalidParamError: If weights are missing or have an unsupported shape.
+        """
+        supported_size = [4]
+        weight_name = node.input[1]
+        initializer = initializer_map.get(weight_name)
+
+        if initializer is None:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"Weight tensor '{weight_name}' not found in initializers",
+            )
+
+        weight_dims = list(initializer.dims)
+
+        if len(weight_dims) not in supported_size:
+            msg = f"Unsupported Conv weight dimensionality {len(weight_dims)}. "
+            msg += f"Expected 4D weights for Conv2D, got shape {weight_dims}"
+            raise InvalidParamError(node.name, node.op_type, msg)

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

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import numpy as np
+import onnx
+from onnx import numpy_helper
+
+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,
+    ScaleConfig,
+)
+
+
+class GemmQuantizer(BaseOpQuantizer):
+    """
+    Quantizer for ONNX Gemm layers.
+
+    - Replaces standard Gemm with Int64Gemm from the `ai.onnx.contrib`
+        domain and makes relevant additional changes to the graph.
+    - Validates that all required Gemm parameters are present.
+    """
+
+    def __init__(
+        self: GemmQuantizer,
+        new_initializers: dict[str, onnx.TensorProto],
+    ) -> None:
+        self.new_initializers = new_initializers
+
+    def quantize(
+        self: GemmQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        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
+
+    def check_supported(
+        self: GemmQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        """
+        Perform high-level validation to ensure that this node
+        can be quantized safely.
+
+        Args:
+            node (onnx.NodeProto): ONNX node to be checked
+            initializer_map (dict[str, onnx.TensorProto]):
+                Initializer map (name of weight or bias and tensor)
+
+        Raises:
+            InvalidParamError: If any requirement is not met.
+        """
+        _ = initializer_map
+        num_valid_inputs = 2
+        # Ensure inputs exist
+        if len(node.input) < num_valid_inputs:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"Expected at least 2 inputs (input, weights), got {len(node.input)}",
+            )
+        num_valid_inputs = 3
+
+        if len(node.input) < num_valid_inputs:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                "Expected at least 3 inputs (input, weights, bias)"
+                f", got {len(node.input)}",
+            )
+
+        # Validate attributes with defaults
+        attrs = {attr.name: attr for attr in node.attribute}
+        alpha = getattr(attrs.get("alpha"), "f", 1.0)
+        beta = getattr(attrs.get("beta"), "f", 1.0)
+        trans_a = getattr(attrs.get("transA"), "i", 0)
+        trans_b = getattr(attrs.get("transB"), "i", 1)
+
+        if alpha != 1.0:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"alpha value of {alpha} not supported",
+                "alpha",
+                "1.0",
+            )
+        if beta != 1.0:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"beta value of {beta} not supported",
+                "beta",
+                "1.0",
+            )
+        if trans_a not in [0, 1]:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"transA value of {trans_a} not supported",
+                "transA",
+                "(0,1)",
+            )
+        if trans_b not in [0, 1]:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"transB value of {trans_b} not supported",
+                "transB",
+                "(0,1)",
+            )

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

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import onnx
+from onnx import helper
+
+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,
+    ScaleConfig,
+)
+
+
+class MaxpoolQuantizer(BaseOpQuantizer):
+    """
+    Quantizer for ONNX MaxPool layers.
+
+    - Replaces standard MaxPool with Int64MaxPool from the `ai.onnx.contrib`
+        domain and makes relevant additional changes to the graph.
+    - Validates that all required MaxPool parameters are present.
+    """
+
+    def __init__(
+        self: MaxpoolQuantizer,
+        new_initializer: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        self.accepted_kernel_shapes = [2]
+        _ = new_initializer
+
+    def quantize(
+        self: BaseOpQuantizer,
+        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,
+            ),
+        ]
+
+    def check_supported(
+        self: MaxpoolQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> None:
+        """
+        Perform high-level validation to ensure that this node
+        can be quantized safely.
+
+        Args:
+            node (onnx.NodeProto): ONNX node to be checked
+            initializer_map (dict[str, onnx.TensorProto]):
+                Initializer map (name of weight or bias and tensor)
+
+        Raises:
+            InvalidParamError: If any requirement is not met.
+        """
+        _ = initializer_map
+        self.check_all_params_exist(node)
+        self.check_params_size(node)
+
+    def check_all_params_exist(self: MaxpoolQuantizer, node: onnx.NodeProto) -> None:
+        """Checks all parameters that are needed, do exist
+
+        Args:
+            node (onnx.NodeProto): ONNX node to check
+
+        Raises:
+            InvalidParamError: If shape requirement is not met.
+        """
+        # May need: ["strides", "kernel_shape", "pads", "dilations"]
+        required_attrs = ["strides", "kernel_shape"]
+        self.validate_required_attrs(node, required_attrs)
+
+        # Check dimension of kernel
+        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,
+                "Currently only MaxPool2D is supported."
+                f"Found {len(kernel_shape)}D kernel",
+                "kernel_shape",
+                "2D",
+            )
+
+    def check_params_size(self: MaxpoolQuantizer, node: onnx.NodeProto) -> None:
+        """Checks dimension of the layer and ensures that it is supported
+
+        Args:
+            node (onnx.NodeProto): ONNX node to check
+
+        Raises:
+            InvalidParamError: If shape requirement is not met.
+        """
+
+        kernel_shape = get_attribute_ints(node, "kernel_shape", default="N/A")
+        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",
+            )

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

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import onnx
+
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    ScaleConfig,
+)
+
+
+class ReluQuantizer(BaseOpQuantizer):
+    """
+    Quantizer for ONNX ReLU layers.
+
+    - Replaces standard ReLU with Int64ReLU from the `ai.onnx.contrib` domain
+        and makes relevant additional changes to the graph.
+    - Validates that all required ReLU parameters are present.
+    """
+
+    def __init__(
+        self: ReluQuantizer,
+        new_initializer: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        _ = new_initializer
+
+    def quantize(
+        self: ReluQuantizer,
+        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]: 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",
+            ),
+        ]
+
+    def check_supported(
+        self: ReluQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        """
+        Perform high-level validation to ensure that this node
+        can be quantized safely.
+
+        Args:
+            node (onnx.NodeProto): ONNX node to be checked
+            initializer_map (dict[str, onnx.TensorProto]):
+                Initializer map (name of weight or bias and tensor)
+        """
+        _ = node
+        _ = initializer_map