JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl

python/core/model_processing/onnx_custom_ops/relu.py

@@ -0,0 +1,43 @@
+import numpy as np
+from onnxruntime_extensions import PyCustomOpDef, onnx_op
+
+
+@onnx_op(
+    op_type="Int64Relu",
+    domain="ai.onnx.contrib",
+    inputs=[PyCustomOpDef.dt_int64],
+    outputs=[PyCustomOpDef.dt_int64],
+)
+def int64_relu(x: np.ndarray) -> np.ndarray:
+    """
+    Performs a ReLU operation on int64 input tensors.
+
+    This function is registered as a custom ONNX operator via onnxruntime_extensions
+    and is used in the JSTprove quantized inference pipeline.
+    It applies ReLU as is (there are no attributes to ReLU).
+
+    Parameters
+    ----------
+    X : Input tensor with dtype int64.
+
+    Returns
+    -------
+    numpy.ndarray
+        ReLU tensor with dtype int64.
+
+    Notes
+    -----
+    - This op is part of the `ai.onnx.contrib` custom domain.
+    - ONNX Runtime Extensions is required to register this op.
+
+    References
+    ----------
+    For more information on the ReLU operation, please refer to the
+    ONNX standard ReLU operator documentation:
+    https://onnx.ai/onnx/operators/onnx__Relu.html
+    """
+    try:
+        return np.maximum(x, 0).astype(np.int64)
+    except Exception as e:
+        msg = f"Int64Gemm failed: {e}"
+        raise RuntimeError(msg) from e

python/core/model_processing/onnx_quantizer/exceptions.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+REPORTING_URL = " https://discord.com/invite/inferencelabs"
+
+
+class QuantizationError(Exception):
+    """
+    Base exception class for errors raised during model quantization.
+    Can be extended for specific quantization-related errors.
+    """
+
+    GENERIC_MESSAGE = (
+        "\nThis model is not supported by JSTprove."
+        f"\nSubmit model support requests via the JSTprove channel: {REPORTING_URL},"
+    )
+
+    def __init__(self: Exception, message: str) -> None:
+        """Initialize QuantizationError with a detailed message.
+
+        Args:
+            message (str): Specific error message describing the quantization issue.
+        """
+        full_msg = f"{self.GENERIC_MESSAGE}\n\n{message}"
+        super().__init__(full_msg)
+
+
+class InvalidParamError(QuantizationError):
+    """
+    Exception raised when invalid parameters or unsupported
+    parameters are encountered in a node that is reached during
+    quantization the quantization process.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self: QuantizationError,
+        node_name: str,
+        op_type: str,
+        message: str,
+        attr_key: str | None = None,
+        expected: str | None = None,
+    ) -> None:
+        """Initialize InvalidParamError with context about the invalid parameter.
+
+        Args:
+            node_name (str): The name of the node where the error occurred.
+            op_type (str): The type of operation of the node.
+            message (str): Description of the invalid parameter error.
+            attr_key (str, optional): The attribute key that caused the error.
+                Defaults to None.
+            expected (str, optional): The expected value or format for the attribute.
+                Defaults to None.
+        """
+        self.node_name = node_name
+        self.op_type = op_type
+        self.message = message
+        self.attr_key = attr_key
+        self.expected = expected
+
+        error_msg = (
+            f"Invalid parameters in node '{node_name}' "
+            f"(op_type='{op_type}'): {message}"
+        )
+        if attr_key:
+            error_msg += f" [Attribute: {attr_key}]"
+        if expected:
+            error_msg += f" [Expected: {expected}]"
+        super().__init__(error_msg)
+
+
+class UnsupportedOpError(QuantizationError):
+    """
+    Exception to be raised when an unsupported operation type is
+    reached during quantization.
+    """
+
+    def __init__(
+        self: QuantizationError,
+        op_type: str,
+        node_name: str | None = None,
+    ) -> None:
+        """Initialize UnsupportedOpError with details about the unsupported operation.
+
+        Args:
+            op_type (str): The type of the unsupported operation.
+            node_name (str, optional): The name of the node where the
+                unsupported operation was found to help with debugging.
+                    Defaults to None.
+        """
+        error_msg = f"Unsupported op type: '{op_type}'"
+        if node_name:
+            error_msg += f" in node '{node_name}'"
+        error_msg += ". Please check out the documentation for supported layers."
+        self.unsupported_ops = op_type
+        super().__init__(error_msg)
+
+
+class MissingHandlerError(QuantizationError):
+    """
+    Raised when no handler is registered for an operator.
+    """
+
+    def __init__(self: QuantizationError, op_type: str) -> None:
+        error_msg = f"No quantization handler registered for operator type '{op_type}'."
+        super().__init__(error_msg)
+
+
+class InitializerNotFoundError(QuantizationError):
+    """
+    Raised when an initializer required by a node is missing from the initializer map.
+    """
+
+    def __init__(
+        self: QuantizationError,
+        node_name: str,
+        initializer_name: str,
+    ) -> None:
+        error_msg = (
+            f"Initializer '{initializer_name}' required for node '{node_name}' "
+            "was not found in the initializer map."
+        )
+        super().__init__(error_msg)
+
+
+class HandlerImplementationError(QuantizationError):
+    """
+    Raised when a handler does not conform to the expected quantizer interface.
+    For example, missing 'quantize' method, wrong return type, etc.
+    """
+
+    def __init__(self: QuantizationError, op_type: str, message: str) -> None:
+        error_msg = f"Handler implementation error for operator '{op_type}': {message}"
+        super().__init__(error_msg)
+
+
+class InvalidGraphError(QuantizationError):
+    """
+    Raised when the ONNX graph is malformed or missing critical information.
+    """
+
+    def __init__(self: QuantizationError, message: str) -> None:
+        error_msg = f"Invalid ONNX graph structure: {message}"
+        super().__init__(error_msg)
+
+
+class InvalidConfigError(QuantizationError):
+    """
+    Exception raised when the overall quantization configuration is invalid
+    or unsupported (e.g., bad scale_base, scale_exponent, or global settings).
+    """
+
+    def __init__(
+        self: QuantizationError,
+        key: str,
+        value: str | float | bool | None,
+        expected: str | None = None,
+    ) -> None:
+        """Initialize InvalidConfigError with context about the bad config.
+
+        Args:
+            key (str): The name of the configuration parameter.
+            value (Union[str, int, float, bool, None]):
+                The invalid value that was provided.
+            expected (str, optional): Description of the expected valid range or type.
+        """
+        error_msg = f"Invalid configuration for '{key}': got {value}"
+        if expected:
+            error_msg += f" (expected {expected})"
+        super().__init__(error_msg)

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

@@ -0,0 +1,396 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+
+from python.core.model_processing.onnx_custom_ops.onnx_helpers import (
+    replace_input_references,
+)
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    HandlerImplementationError,
+    InitializerNotFoundError,
+    InvalidConfigError,
+    InvalidParamError,
+)
+
+
+@dataclass
+class ScaleConfig:
+    exponent: int
+    base: int
+    rescale: bool
+
+
+class BaseOpQuantizer:
+    """
+    Abstract base class for ONNX operator quantizers.
+
+    Subclasses must implement:
+        - `quantize`: Apply quantization logic to an ONNX node.
+        - `check_supported`: Checks if the layer and param specs are supported.
+
+    Attributes:
+        new_initializers (list[onnx.TensorProto]):
+            A list of initializers created during quantization.
+            These should be added to the graph after processing.
+    """
+
+    def __init__(self: BaseOpQuantizer) -> None:
+        self.new_initializers: list[onnx.TensorProto] = []
+
+    @staticmethod
+    def get_scaling(scale_base: int, scale_exponent: int) -> int:
+        """Validate and compute the scaling factor.
+
+        Args:
+            scale_base (int): Base for the scaling exponent.
+            scale_exponent (int): Scaling exponent.
+
+        Returns:
+            int: The computed scaling factor (scale_base ** scale_exponent).
+
+        Raises:
+            InvalidConfigError: If parameters are invalid.
+        """
+        if scale_base <= 0:
+            key = "scale_base"
+            raise InvalidConfigError(key, scale_base, expected="> 0")
+        if scale_exponent < 0:
+            key = "scale_exponent"
+            raise InvalidConfigError(key, scale_exponent, expected=">= 0")
+
+        try:
+            return scale_base**scale_exponent
+        except (TypeError, OverflowError, ValueError, Exception) as e:
+            key = "scaling"
+            raise InvalidConfigError(
+                key,
+                f"{scale_base}^{scale_exponent}",
+                str(e),
+            ) from e
+
+    @staticmethod
+    def validate_node_has_output(node: onnx.NodeProto) -> None:
+        """Ensure a node has at least one output.
+
+        Args:
+            node (onnx.NodeProto): The node to validate.
+            op_type (str): Name of the operator type for error reporting.
+
+        Raises:
+            HandlerImplementationError: If the node has no outputs.
+        """
+        if not node.output or len(node.output) == 0:
+            raise HandlerImplementationError(
+                op_type=node.op_type,
+                message=f"Node '{node.name or '<unnamed>'}' of type '{node.op_type}'"
+                " has no outputs.",
+            )
+
+    @staticmethod
+    def validate_required_attrs(
+        node: onnx.NodeProto,
+        required_attrs: list[str],
+    ) -> None:
+        """
+        Ensure that a node contains all required attributes.
+
+        Args:
+            node (onnx.NodeProto): The ONNX node to validate.
+            required_attrs (list[str]): list of attribute names that must exist.
+            op_type (str): Name of the operator type for error reporting.
+
+        Raises:
+            InvalidParamError: If any required attribute is missing.
+        """
+        missing_attrs = []
+        for attr_name in required_attrs:
+            found = any(attr.name == attr_name for attr in node.attribute)
+            if not found:
+                missing_attrs.append(attr_name)
+
+        if missing_attrs:
+            missing_str = ", ".join(missing_attrs)
+            raise InvalidParamError(
+                node_name=node.name,
+                op_type=node.op_type,
+                message=f"Missing required attributes: {missing_str}",
+            )
+
+    def quantize(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        """
+        Quantize the given node.
+
+        Must be implemented by subclasses.
+
+        Raises:
+            HandlerImplementationError: If subclass does not implement quantize
+        """
+        _ = node, graph, scale_config, initializer_map
+        raise HandlerImplementationError(
+            op_type=self.__class__.__name__,
+            message="quantize() not implemented in subclass.",
+        )
+
+    def check_supported(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> str | None:
+        """
+        Check if the node is supported by the quantizer.
+
+        Must be overridden by subclasses to validate parameters.
+
+        Raises:
+            HandlerImplementationError: If called on BaseOpQuantizer directly.
+        """
+        _ = node, initializer_map
+        raise HandlerImplementationError(
+            op_type=self.__class__.__name__,
+            message="check_supported() not implemented in subclass.",
+        )
+
+    def rescale_layer(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        scale_base: int,
+        scale_exponent: int,
+        graph: onnx.GraphProto,
+    ) -> list[onnx.NodeProto]:
+        """
+        Helper function for any quantizer.
+        Used to add a rescaling step after the given node.
+
+        This replaces the node's output with a scaled version using a Div op.
+        This function incorporates the logic to insert and restructure the graph.
+
+        Args:
+            node (onnx.NodeProto): Node to rescale.
+            scale_base (int): Base for the scaling exponent.
+            scale_exponent (int): Scaling exponent.
+            graph (onnx.GraphProto): The ONNX graph.
+
+        Returns:
+            list[onnx.NodeProto]: Original node and the inserted Div node.
+
+        Raises:
+            HandlerImplementationError if there are no outputs to be rescaled
+        """
+        self.validate_node_has_output(node)
+
+        original_output = node.output.get(0)
+        quantized_output = original_output + "_raw"
+        node.output[0] = quantized_output
+
+        # Create scale constant initializer
+        scale_const_name = node.name + "_scale"
+
+        scale_value = self.get_scaling(scale_base, scale_exponent)
+        scale_tensor = numpy_helper.from_array(
+            np.array([scale_value], dtype=np.int64),
+            name=scale_const_name,
+        )
+        self.new_initializers.append(scale_tensor)
+
+        # Create Div node for rescaling output
+        div_node = helper.make_node(
+            "Div",
+            inputs=[quantized_output, scale_const_name],
+            outputs=[original_output],  # restore original output name
+            name=node.name + "_rescale",
+        )
+
+        # Rewire consumers to point to the new output
+        replace_input_references(
+            graph=graph,
+            old_output=original_output,
+            new_output=div_node.output[0],
+        )
+
+        return [node, div_node]
+
+    def add_nodes_w_and_b(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        scale_exponent: int,
+        scale_base: int,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> tuple[list[onnx.NodeProto], list[str]]:
+        """Insert scaling and casting nodes for weight and bias,
+            to convert from float to scaled int64 values.
+
+        Args:
+            node (onnx.NodeProto): Node to find used weights and biases.
+            scale_exponent (int): Scaling exponent.
+            scale_base (int): Base for the scaling exponent.
+            initializer_map (dict[str, onnx.TensorProto]): The initializer map.
+            graph (onnx.GraphProto): ONNX Graph
+
+        Returns:
+            tuple[list[onnx.NodeProto], list[str]]:
+                list of new nodes added, updated input names for nodes.
+
+        Raises:
+            InitializerNotFoundError: If weights or biases are missing from the graph.
+            HandlerImplementationError:
+                If there are no weights or biases to add to the graph.
+        """
+        weights_input_length = 2
+        if len(node.input) < weights_input_length:
+            raise HandlerImplementationError(
+                op_type=node.op_type,
+                message=f"Node '{node.name}'"
+                " has fewer than 2 inputs (weights missing).",
+            )
+        # Quantize weight
+        weight_name = node.input[1]
+        if not weight_name or weight_name not in initializer_map:
+            raise InitializerNotFoundError(node.name, weight_name or "<missing>")
+
+        weight_tensor = initializer_map[weight_name]
+        if not weight_tensor.name:
+            raise HandlerImplementationError(
+                op_type=node.op_type,
+                message=f"Weight tensor for node '{node.name}' is missing a name.",
+            )
+
+        quant_weight_name, mul_node, cast_node = self.insert_scale_node(
+            tensor=weight_tensor,
+            scale_base=scale_base,
+            scale_exponent=scale_exponent,
+        )
+
+        # Quantize bias if present
+        new_inputs = [node.input[0], quant_weight_name]
+        nodes = [mul_node, cast_node]
+
+        bias_inputs_length = 3
+
+        if len(node.input) >= bias_inputs_length:
+            bias_name = node.input[2]
+            if bias_name not in initializer_map:
+                raise InitializerNotFoundError(node.name, bias_name)
+
+            bias_tensor = initializer_map[bias_name]
+            quant_bias_name, mul_node_2, cast_node_2 = self.insert_scale_node(
+                tensor=bias_tensor,
+                scale_base=scale_base,
+                scale_exponent=(scale_exponent * 2),
+            )
+            new_inputs.append(quant_bias_name)
+            nodes.append(mul_node_2)
+            nodes.append(cast_node_2)
+
+        # === Mutate the original node ===
+        return nodes, new_inputs
+
+    def insert_scale_node(
+        self: BaseOpQuantizer,
+        tensor: onnx.TensorProto,
+        scale_base: int,
+        scale_exponent: int,
+    ) -> tuple[str, onnx.NodeProto, onnx.NodeProto]:
+        """Insert Mul and Cast nodes to apply scaling to a tensor.
+
+        Args:
+            tensor (onnx.TensorProto): Tensor to scale.
+            scale_base (int): Base for scaling exponent.
+            scale_exponent (int): Scaling exponent.
+            graph (onnx.GraphProto): ONNX graph.
+
+        Returns:
+            tuple[str, onnx.NodeProto, onnx.NodeProto]:
+                New tensor name, Mul node, Cast node.
+
+        Raises:
+            HandlerImplementationError:
+                If tensor does not exist, incorrectly formatted or not named
+        """
+        if not tensor or not isinstance(tensor, onnx.TensorProto):
+            raise HandlerImplementationError(
+                op_type="insert_scale_node",
+                message="Expected a valid onnx.TensorProto, got None or wrong type.",
+            )
+
+        if not tensor.name:
+            raise HandlerImplementationError(
+                op_type="insert_scale_node",
+                message="Tensor is missing a name.",
+            )
+
+        scale_value = self.get_scaling(scale_base, scale_exponent)
+
+        # Create scale constant
+        scale_const_name = tensor.name + "_scale"
+        scale_tensor = numpy_helper.from_array(
+            np.array([scale_value], dtype=np.float64),
+            name=scale_const_name,
+        )
+        self.new_initializers.append(scale_tensor)
+
+        # Add Mul node
+        scaled_output_name = f"{tensor.name}_scaled"
+        mul_node = helper.make_node(
+            "Mul",
+            inputs=[tensor.name, scale_const_name],
+            outputs=[scaled_output_name],
+            name=f"{tensor.name}_mul",
+        )
+
+        # Add cast node
+        output_name = f"{scaled_output_name}_cast"
+        rounded_output_name = scaled_output_name
+        cast_to_int64 = helper.make_node(
+            "Cast",
+            inputs=[scaled_output_name],
+            outputs=[output_name],
+            to=onnx.TensorProto.INT64,
+            name=rounded_output_name,
+        )
+        return output_name, mul_node, cast_to_int64
+
+
+class PassthroughQuantizer(BaseOpQuantizer):
+    """
+    Quantizer that leaves the node unchanged.
+    Useful for operators that do not require quantization, such as shaping operations.
+    """
+
+    def __init__(
+        self: BaseOpQuantizer,
+        new_initializer: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        _ = new_initializer
+        super().__init__()
+
+    def quantize(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        _ = graph, scale_config, initializer_map
+        if not isinstance(node, onnx.NodeProto):
+            raise HandlerImplementationError(
+                op_type="PassthroughQuantizer",
+                message="quantize() expected a NodeProto",
+            )
+        return [node]
+
+    def check_supported(
+        self: BaseOpQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        _ = node, initializer_map