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

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

@@ -418,6 +418,40 @@ class BaseOpQuantizer:
 
 
 class QuantizerBase:
+    """
+    Shared mixin implementing the generic INT64 quantization pipeline.
+
+    IMPORTANT:
+        QuantizerBase is *not* a standalone quantizer. It must always be
+        combined with BaseOpQuantizer via multiple inheritance:
+
+            class FooQuantizer(BaseOpQuantizer, QuantizeFoo):
+                ...
+
+        BaseOpQuantizer supplies required methods and attributes that
+        QuantizerBase relies on:
+            - add_scaled_initializer_inputs
+            - insert_scale_node
+            - get_scaling
+            - new_initializers  (initializer buffer shared with converter)
+
+        If a subclass inherits QuantizerBase without BaseOpQuantizer,
+        QuantizerBase.quantize() will raise attribute errors at runtime.
+
+    This mixin centralizes:
+        - attribute extraction/merging
+        - optional initializer scaling (USE_WB + SCALE_PLAN)
+        - optional rescaling of outputs (USE_SCALING)
+        - creation of the final quantized NodeProto
+
+    The Quantize<Op> mixins should define:
+        - OP_TYPE
+        - DOMAIN
+        - USE_WB (bool)
+        - USE_SCALING (bool)
+        - SCALE_PLAN (dict[int,int]) if initializer scaling is enabled
+    """
+
     OP_TYPE = None
     DOMAIN = "ai.onnx.contrib"
     DEFAULT_ATTRS: ClassVar = {}
@@ -479,6 +513,73 @@ class QuantizerBase:
         nodes.append(quantized_node)
         return nodes
 
+    def pre_analysis_transform(
+        self: QuantizerBase,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        initializer_map: dict[str, onnx.TensorProto],
+        scale_base: int,
+        scale_exponent: int,
+    ) -> None:
+        """
+        pre_analysis_transform aims to transform the given layer along the
+        same lines as it would be transformed for the quantized model, but
+        for the weights and biases file instead, to be sent to the backend
+
+        Default pre-analysis behavior:
+
+        - If the subclass uses weights/bias (`USE_WB=True`), apply the SAME
+        scaling rules as quantization, but directly mutate the initializers.
+
+        - Subclasses can override this to implement more complex rewrites
+        (e.g., BatchNorm → Mul/Add).
+
+        Args:
+            node (onnx.NodeProto): Node to transform.
+            graph (onnx.GraphProto): Rest of the Onnx graph for initializers.
+            initializer_map (dict[str, onnx.TensorProto]): The initializer map.
+
+            scale_base (int): Scaling base.
+            scale_exponent (int): Scaling exponent.
+
+        NOTE
+         - The resulting model will not make accurate prediction and should be
+         used solely for analysis and keeping track of w_and_b
+        """
+        # If subclass does not want auto-scaling, do nothing
+        if not getattr(self, "USE_WB", False):
+            return
+
+        # Each quantizer defines which inputs to scale (Weight:1x, Bias:2x etc.)
+        scale_plan = getattr(self, "SCALE_PLAN", {})
+
+        # Perform the same scaling as quantization, but directly modify initializers
+        for input_idx, scale_mult in scale_plan.items():
+            if input_idx >= len(node.input):
+                continue
+
+            name = node.input[input_idx]
+            if name not in initializer_map:
+                continue  # optional input missing
+
+            tensor = initializer_map[name]
+            arr = numpy_helper.to_array(tensor).astype(np.float64)
+
+            scale = scale_base ** (scale_exponent * scale_mult)
+            new_arr = arr * scale
+
+            # Replace initializer directly
+            new_tensor = numpy_helper.from_array(new_arr, name=tensor.name)
+
+            # Modify graph initializer in place
+            for j in range(len(graph.initializer)):
+                if graph.initializer[j].name == tensor.name:
+                    del graph.initializer[j]
+                    break
+            graph.initializer.append(new_tensor)
+
+            initializer_map[tensor.name] = new_tensor
+
 
 class PassthroughQuantizer(BaseOpQuantizer):
     """

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

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+from python.core.circuits.errors import CircuitConfigurationError
+
+if TYPE_CHECKING:
+    import onnx
+
+import numpy as np
+from onnx import helper, 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,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeBatchnorm(QuantizerBase):
+    OP_TYPE = "Int64BatchNorm"
+    USE_WB = True
+    USE_SCALING = False
+    SCALE_PLAN: ClassVar = {}
+
+
+class BatchnormQuantizer(BaseOpQuantizer, QuantizeBatchnorm):
+    """
+    Quantizer for ONNX Batchnorm layers.
+
+    - Uses standard ONNX Batchnorm layer in standard domain, and
+      makes relevant additional changes to the graph.
+    """
+
+    def __init__(
+        self: BatchnormQuantizer,
+        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 _compute_mul_add(
+        self: BatchnormQuantizer,
+        initializer_map: dict[str, onnx.TensorProto],
+        node: onnx.NodeProto,
+        scale_base: int,
+        scale_exponent: int,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Compute the 'mul' and 'add' tensors for BatchNorm folding.
+        """
+        self._validate_inputs(node=node)
+        # ONNX BatchNorm inputs: [X, scale, bias, mean, var]
+        scale_factor = scale_base**scale_exponent
+        scale = numpy_helper.to_array(initializer_map[node.input[1]]).astype(np.float32)
+        bias = numpy_helper.to_array(initializer_map[node.input[2]]).astype(np.float32)
+        mean = numpy_helper.to_array(initializer_map[node.input[3]]).astype(np.float32)
+        var = numpy_helper.to_array(initializer_map[node.input[4]]).astype(np.float32)
+
+        # Find epsilon attribute
+        epsilon_attr = next((a for a in node.attribute if a.name == "epsilon"), None)
+        epsilon = float(epsilon_attr.f) if epsilon_attr else 1e-5
+
+        mul = scale / np.sqrt(var + epsilon)
+        add = bias - mean * mul
+        scaled_add = add * (scale_factor**2)
+        scaled_mul = scale_factor * mul
+        return scaled_mul, scaled_add
+
+    def pre_analysis_transform(
+        self: BatchnormQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        initializer_map: dict[str, onnx.TensorProto],
+        scale_base: int,
+        scale_exponent: int,
+    ) -> None:
+        # Compute linearized BN tensors
+        mul, add = self._compute_mul_add(
+            initializer_map,
+            node,
+            scale_base=scale_base,
+            scale_exponent=scale_exponent,
+        )
+
+        # Name base
+        node_name = node.name if node.name else node.input[0]
+        mul_name = f"{node_name}_mul"
+        add_name = f"{node_name}_add"
+
+        # Create ONNX tensors
+        mul_tensor = numpy_helper.from_array(mul.astype(np.int64), name=mul_name)
+        add_tensor = numpy_helper.from_array(add.astype(np.int64), name=add_name)
+
+        # Insert them into the graph
+        graph.initializer.extend([mul_tensor, add_tensor])
+        initializer_map[mul_name] = mul_tensor
+        initializer_map[add_name] = add_tensor
+        self.new_initializers.extend([mul_tensor, add_tensor])
+
+        node.input[:] = [node.input[0], mul_name, add_name]
+
+        del node.attribute[:]
+
+    def quantize(
+        self,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        _ = graph
+
+        nodes: list[onnx.NodeProto] = []
+
+        # 1. Compute unscaled float mul/add coefficients
+        mul, add = self._compute_mul_add(
+            initializer_map,
+            node,
+            scale_base=1,
+            scale_exponent=1,
+        )
+
+        node_name = node.name if node.name else node.input[0]
+        mul_name = f"{node_name}_mul"
+        add_name = f"{node_name}_add"
+
+        # 2. Store unscaled mul and add initializers (as floats)
+        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)
+
+        mul_tensor = numpy_helper.from_array(mul.astype(np.float32), name=mul_name)
+        add_tensor = numpy_helper.from_array(add.astype(np.float32), name=add_name)
+
+        initializer_map[mul_name] = mul_tensor
+        initializer_map[add_name] = add_tensor
+
+        # 3. Insert scale and cast for mul_tensor
+        scaled_mul_name, mul_scale_node, mul_cast_node = self.insert_scale_node(
+            tensor=mul_tensor,
+            scale_base=scale_config.base,
+            scale_exponent=scale_config.exponent,
+        )
+
+        # 4. Insert scale and cast for add_tensor
+        scaled_add_name, add_scale_node, add_cast_node = self.insert_scale_node(
+            tensor=add_tensor,
+            scale_base=scale_config.base,
+            scale_exponent=scale_config.exponent * 2,
+        )
+        # Note, order is important here
+        nodes.extend(
+            [
+                mul_scale_node,
+                mul_cast_node,
+                add_scale_node,
+                add_cast_node,
+            ],
+        )
+
+        # 5. Build final Int64BatchNorm node
+        attrs = extract_attributes(node)
+        for k, v in getattr(self, "DEFAULT_ATTRS", {}).items():
+            attrs.setdefault(k, v)
+        attrs["rescale"] = 1
+
+        quant_node = helper.make_node(
+            self.OP_TYPE,  # Should be "Int64BatchNorm"
+            inputs=[
+                node.input[0],  # original X
+                scaled_mul_name,  # scaled mul
+                scaled_add_name,  # scaled add
+                scale_name,  # scaling factor
+            ],
+            outputs=node.output,
+            name=node.name,
+            domain=self.DOMAIN,
+            **attrs,
+        )
+
+        nodes.append(quant_node)
+        return nodes
+
+    def check_supported(
+        self: BatchnormQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        """
+        For our current implementation, all batchnorm inputs
+        (scale, variance, mean, etc.)
+        must be initializers to the circuit and not inputs from earlier in the graph.
+        """
+
+        if initializer_map is None:
+            msg = "initializer_map is required for BatchNorm support check"
+            raise CircuitConfigurationError(node.name, node.op_type, msg)
+
+        self._validate_inputs(node=node)
+
+        # First, check to make sure that each of the batchnorm inputs are initializers
+        initializer_inputs = node.input[1:]
+        if not all(i in initializer_map for i in initializer_inputs):
+            msg = "Unsupported BatchNorm with normalization inputs not in initializers"
+            raise InvalidParamError(node.name, node.op_type, msg)
+
+    def _validate_inputs(self, node: onnx.NodeProto) -> None:
+        """Validate BatchNorm has required inputs in initializer_map."""
+        num_inputs = 5
+        if len(node.input) < num_inputs:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"BatchNorm requires 5 inputs, got {len(node.input)}",
+            )

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

@@ -0,0 +1,92 @@
+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 QuantizeClip(QuantizerBase):
+    """
+    Quantization traits for ONNX Clip.
+
+    Semantics:
+    - X is already scaled/cast to INT64 at the graph boundary by the converter.
+    - Clip is elementwise + broadcasting.
+    - The bound inputs (min, max) should live in the *same* fixed-point scale
+      as X so that Clip(alpha*x; alpha*a, alpha*b) matches the original Clip(x; a, b).
+
+    Implementation:
+    - Treat inputs 1 and 2 (min, max) like "WB-style" slots: we let the
+      QuantizerBase machinery rescale / cast those inputs using the same
+      global scale factor.
+    - No extra internal scaling input is added (USE_SCALING = False).
+    """
+
+    OP_TYPE = "Clip"
+    DOMAIN = ""  # standard ONNX domain
+
+    # We DO want WB-style handling so that min/max initializers get quantized:
+    USE_WB = True
+
+    # Clip does not introduce its own scale input; it just runs in the
+    # existing fixed-point scale.
+    USE_SCALING = False
+
+    # Scale-plan for WB-style slots:
+    #   - Input index 1: min
+    #   - Input index 2: max
+    # Each should be scaled once by the global alpha (same as activations).
+    SCALE_PLAN: ClassVar = {1: 1, 2: 1}
+
+
+class ClipQuantizer(BaseOpQuantizer, QuantizeClip):
+    """
+    Quantizer for ONNX Clip.
+
+    - Keeps the node op_type as "Clip".
+    - Ensures that any bound inputs (min, max), whether they are dynamic
+      inputs or initializers, are converted to the same INT64 fixed-point
+      representation as A.
+    """
+
+    def __init__(
+        self,
+        new_initializers: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        # Match Max/Min/Add: we simply share the new_initializers dict
+        # with the converter so any constants we add are collected.
+        self.new_initializers = new_initializers
+
+    def quantize(
+        self,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        # Delegate to the shared QuantizerBase logic, which will:
+        # - keep X as-is (already scaled/cast by the converter),
+        # - rescale / cast min/max according to SCALE_PLAN,
+        # - update initializers as needed.
+        return QuantizeClip.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        """
+        Minimal support check for Clip:
+
+        - Clip is variadic elementwise with optional min/max as inputs or attrs.
+        - We accept both forms; if attrs are present, ORT enforces semantics.
+        - Broadcasting is ONNX-standard; we don't restrict further here.
+        """
+        _ = node, initializer_map

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

@@ -0,0 +1,49 @@
+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 QuantizeMax(QuantizerBase):
+    OP_TYPE = "Max"
+    DOMAIN = ""
+    USE_WB = True
+    USE_SCALING = False
+    SCALE_PLAN: ClassVar = {1: 1}
+
+
+class MaxQuantizer(BaseOpQuantizer, QuantizeMax):
+    def __init__(
+        self,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        if new_initializers is not None:
+            # Share the caller-provided buffer instead of the default list.
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        # Delegate to the shared QuantizerBase logic
+        return QuantizeMax.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        # If later we want to enforce/relax broadcasting, add it here.
+        pass

python/core/model_processing/onnx_quantizer/layers/min.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 QuantizeMin(QuantizerBase):
+    OP_TYPE = "Min"
+    DOMAIN = ""  # standard ONNX domain
+    USE_WB = True  # let framework wire inputs/outputs normally
+    USE_SCALING = False  # passthrough: no internal scaling
+    SCALE_PLAN: ClassVar = {1: 1}  # elementwise arity plan
+
+
+class MinQuantizer(BaseOpQuantizer, QuantizeMin):
+    """
+    Passthrough quantizer for elementwise Min.
+    We rely on the converter to quantize graph inputs; no extra scaling here.
+    """
+
+    def __init__(
+        self: MinQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: MinQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        # Delegate to QuantizerBase's generic passthrough implementation.
+        return QuantizeMin.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: MinQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        # Min has no attributes; elementwise, variadic ≥ 1 input per ONNX spec.
+        # We mirror Add/Max broadcasting behavior; no extra checks here.
+        _ = node, initializer_map

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