JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl

python/core/model_processing/errors.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from python.core.model_processing.converters.base import ModelType
+
+
+class ModelConversionError(Exception):
+    """Base class for all model conversion errors."""
+
+    def __init__(
+        self: ModelConversionError,
+        message: str,
+        model_type: ModelType,
+        context: dict | None = None,
+    ) -> None:
+        self.message = message
+        self.model_type = model_type
+        self.context = context or {}
+        super().__init__(self.__str__())
+
+    def __str__(self: ModelConversionError) -> str:
+        msg = f"Error converting {self.model_type.value} model: {self.message}"
+        if self.context:
+            ctx_str = ", ".join(f"{k}={v}" for k, v in self.context.items())
+            msg += f"(Context: {ctx_str})"
+        return msg
+
+
+class ModelLoadError(ModelConversionError):
+    """Raised when an ONNX model cannot be loaded."""
+
+    def __init__(
+        self: ModelLoadError,
+        file_path: str,
+        model_type: ModelType,
+        reason: str = "",
+    ) -> None:
+        message = f"Failed to load {model_type.value} model from '{file_path}'."
+        if reason:
+            message += f" Reason: {reason}"
+        super().__init__(message, model_type, context={"file_path": file_path})
+
+
+class ModelSaveError(ModelConversionError):
+    """Raised when saving a model fails."""
+
+    def __init__(
+        self: ModelSaveError,
+        file_path: str,
+        model_type: ModelType,
+        reason: str = "",
+    ) -> None:
+        message = f"Failed to save {model_type.value} model to '{file_path}'."
+        if reason:
+            message += f" Reason: {reason}"
+        super().__init__(message, model_type, context={"file_path": file_path})
+
+
+class InferenceError(ModelConversionError):
+    """Raised when inference via ONNX Runtime fails."""
+
+    def __init__(
+        self: InferenceError,
+        model_type: ModelType,
+        model_path: str | None = None,
+        reason: str = "",
+    ) -> None:
+        message = f" {model_type.value} inference failed."
+        if model_path:
+            message += f" Model: '{model_path}'."
+        if reason:
+            message += f" Reason: {reason}"
+        super().__init__(message, model_type, context={"model_path": model_path})
+
+
+class LayerAnalysisError(ModelConversionError):
+    """Raised when analyzing model layers fails."""
+
+    def __init__(
+        self: LayerAnalysisError,
+        model_type: ModelType,
+        layer_name: str | None = None,
+        reason: str = "",
+    ) -> None:
+        message = "Layer analysis failed."
+        if layer_name:
+            message += f" Problematic layer: '{layer_name}'."
+        if reason:
+            message += f" Reason: {reason}"
+        super().__init__(message, model_type, context={"layer_name": layer_name})
+
+
+class IOInfoExtractionError(ModelConversionError):
+    """Raised when extracting input/output info fails."""
+
+    def __init__(
+        self: IOInfoExtractionError,
+        model_type: ModelType,
+        model_path: str | None = None,
+        reason: str = "",
+    ) -> None:
+        message = "Failed to extract input/output info from model."
+        if model_path:
+            message += f" Model: '{model_path}'."
+        if reason:
+            message += f" Reason: {reason}"
+        super().__init__(message, model_type, context={"model_path": model_path})
+
+
+class InvalidModelError(ModelConversionError):
+    """Raised when an ONNX model fails validation checks (onnx.checker)."""
+
+    def __init__(
+        self: InvalidModelError,
+        model_type: ModelType,
+        model_path: str | None = None,
+        reason: str = "",
+    ) -> None:
+        msg = f"The {model_type.value} model is invalid."
+        if model_path:
+            msg += f" Model: '{model_path}'."
+        if reason:
+            msg += f" Reason: {reason}"
+        super().__init__(
+            message=msg,
+            model_type=model_type,
+            context={"model_path": model_path},
+        )
+
+
+class SerializationError(ModelConversionError):
+    """Raised when model data cannot be serialized to the required format."""
+
+    def __init__(
+        self: SerializationError,
+        model_type: ModelType,
+        tensor_name: str | None = None,
+        reason: str = "",
+    ) -> None:
+        message = "Failed to serialize model data."
+        if tensor_name:
+            message += f" Tensor: '{tensor_name}'."
+        if reason:
+            message += f" Reason: {reason}"
+        super().__init__(message, model_type, context={"tensor_name": tensor_name})

python/core/model_processing/onnx_custom_ops/__init__.py

@@ -0,0 +1,16 @@
+import importlib
+import pkgutil
+import os
+
+# Get the package name of the current module
+package_name = __name__
+
+# Dynamically import all .py files in this package directory (except __init__.py)
+package_dir = os.path.dirname(__file__)
+
+__all__ = []
+
+for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
+    if not is_pkg and (module_name != "custom_helpers"):
+        importlib.import_module(f"{package_name}.{module_name}")
+        __all__.append(module_name)

python/core/model_processing/onnx_custom_ops/conv.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as f
+from onnxruntime_extensions import PyCustomOpDef, onnx_op
+
+from .custom_helpers import parse_attr, rescaling
+
+
+@onnx_op(
+    op_type="Int64Conv",
+    domain="ai.onnx.contrib",
+    inputs=[
+        PyCustomOpDef.dt_int64,  # X
+        PyCustomOpDef.dt_int64,  # W
+        PyCustomOpDef.dt_int64,  # B
+        PyCustomOpDef.dt_int64,  # scaling factor
+    ],
+    outputs=[PyCustomOpDef.dt_int64],
+    attrs={
+        "auto_pad": PyCustomOpDef.dt_string,
+        "strides": PyCustomOpDef.dt_string,
+        "pads": PyCustomOpDef.dt_string,
+        "dilations": PyCustomOpDef.dt_string,
+        "group": PyCustomOpDef.dt_int64,
+        "kernel_shape": PyCustomOpDef.dt_string,
+        "rescale": PyCustomOpDef.dt_int64,
+    },
+)
+def int64_conv(
+    x: np.ndarray,
+    w: np.ndarray,
+    b: np.ndarray | None = None,
+    scaling_factor: np.ndarray | None = None,
+    auto_pad: str | None = None,
+    dilations: str | None = None,
+    group: int | None = None,
+    kernel_shape: str | None = None,
+    pads: str | None = None,
+    strides: str | None = None,
+    rescale: int | None = None,
+) -> np.ndarray:
+    """
+    Performs a convolution 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 parses ONNX-style
+    convolution attributes, applies convolution
+    and optionally rescales the result.
+
+    Parameters
+    ----------
+    X : Input tensor with dtype int64.
+    W : Convolution weight tensor with dtype int64.
+    B : Optional bias tensor with dtype int64.
+    scaling_factor : Scaling factor for rescaling the output.
+    auto_pad : Optional ONNX auto padding type (`SAME_UPPER`, `SAME_LOWER`, `VALID`).
+    dilations : Dilation values for the convolution (default: `[1, 1]`).
+    group : Group value for the convolution (default: 1).
+    kernel_shape : Kernel shape (default: `[3, 3]`).
+    pads : Padding values (default: `[0, 0, 0, 0]`).
+    strides : Stride values (default: `[1, 1]`).
+    rescale : Optional flag to apply output rescaling or not.
+
+    Returns
+    -------
+    numpy.ndarray
+        Convolved 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 convolution operation, please refer to the
+    ONNX standard Conv operator documentation:
+    https://onnx.ai/onnx/operators/onnx__Conv.html
+    """
+    _ = auto_pad
+    try:
+        strides = parse_attr(strides, [1, 1])
+        dilations = parse_attr(dilations, [1, 1])
+        pads = parse_attr(pads, [0, 0, 0, 0])
+        kernel_shape = parse_attr(kernel_shape, [3, 3])
+
+        x = torch.from_numpy(x)
+        w = torch.from_numpy(w)
+        b = torch.from_numpy(b)
+
+        result = (
+            f.conv2d(
+                x,
+                w,
+                bias=b,
+                stride=strides,
+                padding=pads[:2],
+                dilation=dilations,
+                groups=group,
+            )
+            .numpy()
+            .astype(np.int64)
+        )
+        result = rescaling(scaling_factor, rescale, result)
+        return result.astype(np.int64)
+
+    except Exception as e:
+        msg = f"Int64Conv failed: {e}"
+        raise RuntimeError(msg) from e

python/core/model_processing/onnx_custom_ops/custom_helpers.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def rescaling(scaling_factor: int, rescale: int, y: int) -> int:
+    """Applies integer rescaling to a value based on the given scaling factor.
+
+    Args:
+        scaling_factor (int): The divisor to apply when rescaling.
+            Must be provided if `rescale` is True.
+        rescale (int): Whether to apply rescaling. (0 -> no rescaling, 1 -> rescaling).
+        Y (int): The value to be rescaled.
+
+    Raises:
+        NotImplementedError: If `rescale` is 1 but `scaling_factor` is not provided.
+        NotImplementedError: If `rescale` is not 0 or 1.
+
+    Returns:
+        int: The rescaled value if `rescale` is True, otherwise the original value.
+    """
+    if rescale == 1:
+        if scaling_factor is None:
+            msg = "scaling_factor must be specified when rescale=1"
+            raise ValueError(msg)
+        return y // scaling_factor
+    if rescale == 0:
+        return y
+    msg = f"Rescale must be 0 or 1, got {rescale}"
+    raise ValueError(msg)
+
+
+def parse_attr(attr: str, default: T) -> T:
+    """Parses an attribute list of strings into a list of integers.
+
+    Args:
+        attr (str): Attribute to parse. If a string, it must be
+                    comma-separated integers (e.g., "1, 2, 3").
+                    If None, returns `default`.
+        default (T): Default value to return if `attr` is None.
+
+    Raises:
+        ValueError: If `attr` is a string but cannot be parsed into integers.
+
+    Returns:
+        T: Parsed list of integers if attr is provided, otherwise the default value.
+    """
+    if attr is None:
+        return default
+    try:
+        return [int(x.strip()) for x in attr.split(",")]
+    except ValueError as e:
+        msg = f"Invalid attribute format: {attr}"
+        raise ValueError(msg) from e

python/core/model_processing/onnx_custom_ops/gemm.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import numpy as np
+from onnxruntime_extensions import PyCustomOpDef, onnx_op
+
+from .custom_helpers import rescaling
+
+
+@onnx_op(
+    op_type="Int64Gemm",
+    domain="ai.onnx.contrib",
+    inputs=[
+        PyCustomOpDef.dt_int64,  # X
+        PyCustomOpDef.dt_int64,  # W
+        PyCustomOpDef.dt_int64,  # B
+        PyCustomOpDef.dt_int64,  # Scalar
+    ],
+    outputs=[PyCustomOpDef.dt_int64],
+    attrs={
+        "alpha": PyCustomOpDef.dt_float,
+        "beta": PyCustomOpDef.dt_float,
+        "transA": PyCustomOpDef.dt_int64,
+        "transB": PyCustomOpDef.dt_int64,
+        "rescale": PyCustomOpDef.dt_int64,
+    },
+)
+def int64_gemm7(
+    a: np.ndarray,
+    b: np.ndarray,
+    c: np.ndarray | None = None,
+    scaling_factor: np.ndarray | None = None,
+    alpha: float | None = None,
+    beta: float | None = None,
+    transA: int | None = None,  # noqa: N803
+    transB: int | None = None,  # noqa: N803
+    rescale: int | None = None,
+) -> np.ndarray:
+    """
+    Performs a Gemm (alternatively: Linear layer) 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 parses ONNX-style
+    gemm attributes, applies gemm
+    and optionally rescales the result.
+
+    Parameters
+    ----------
+    a : Input tensor with dtype int64.
+    b : Gemm weight tensor with dtype int64.
+    c : Optional bias tensor with dtype int64.
+    scaling_factor : Scaling factor for rescaling the output.
+    alpha : alpha value for Gemm operation.
+    beta : beta value for Gemm operation.
+    transA : Transpose the a matrix before the Gemm operation
+    transB : Transpose the b matrix before the Gemm operation
+    rescale : Optional flag to apply output rescaling or not.
+
+    Returns
+    -------
+    numpy.ndarray
+        Gemm 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 gemm operation, please refer to the
+    ONNX standard Gemm operator documentation:
+    https://onnx.ai/onnx/operators/onnx__Gemm.html
+    """
+    try:
+        alpha = int(alpha)
+        beta = int(beta)
+
+        a = a.T if transA else a
+        b = b.T if transB else b
+
+        result = alpha * (a @ b)
+
+        if c is not None:
+            result += beta * c
+
+        result = rescaling(scaling_factor, rescale, result)
+        return result.astype(np.int64)
+
+    except Exception as e:
+        msg = f"Int64Gemm failed: {e}"
+        raise RuntimeError(msg) from e

python/core/model_processing/onnx_custom_ops/maxpool.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as f
+from onnxruntime_extensions import PyCustomOpDef, onnx_op
+
+from .custom_helpers import parse_attr
+
+
+@onnx_op(
+    op_type="Int64MaxPool",
+    domain="ai.onnx.contrib",
+    inputs=[PyCustomOpDef.dt_int64],  # input tensor
+    outputs=[PyCustomOpDef.dt_int64],
+    attrs={
+        "strides": PyCustomOpDef.dt_string,
+        "pads": PyCustomOpDef.dt_string,
+        "kernel_shape": PyCustomOpDef.dt_string,
+        "dilations": PyCustomOpDef.dt_string,
+    },
+)
+def int64_maxpool(
+    x: np.ndarray,
+    strides: str | None = None,
+    pads: str | None = None,
+    kernel_shape: str | None = None,
+    dilations: str | None = None,
+) -> np.ndarray:
+    """
+    Performs a MaxPool 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 parses ONNX-style
+    maxpool attributes and applies maxpool.
+
+    Parameters
+    ----------
+    X : Input tensor with dtype int64.
+    kernel_shape : Kernel shape (default: `[2, 2]`).
+    pads : Padding values (default: `[0, 0, 0, 0]`).
+    strides : Stride values (default: `[1, 1]`).
+    dilations : dilation values (default: `[1, 1]`).
+
+
+    Returns
+    -------
+    numpy.ndarray
+        Maxpool 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 maxpool operation, please refer to the
+    ONNX standard MaxPool operator documentation:
+    https://onnx.ai/onnx/operators/onnx__MaxPool.html
+    """
+    try:
+        strides = parse_attr(strides, [1, 1])
+        pads = parse_attr(pads, [0, 0])
+        kernel_size = parse_attr(kernel_shape, [2, 2])
+        dilations = parse_attr(dilations, [1, 1])
+
+        x = torch.from_numpy(x)
+        result = f.max_pool2d(
+            x,
+            kernel_size=kernel_size,
+            stride=strides,
+            padding=pads[:2],
+            dilation=dilations,
+        )
+        return result.numpy().astype(np.int64)
+    except Exception as e:
+        msg = f"Int64Gemm failed: {e}"
+        raise RuntimeError(msg) from e

python/core/model_processing/onnx_custom_ops/onnx_helpers.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+from typing import Any
+
+import onnx
+from onnx import AttributeProto, numpy_helper
+
+ATTRIBUTE_PARSERS = {
+    AttributeProto.FLOAT: lambda a: a.f,
+    AttributeProto.INT: lambda a: a.i,
+    AttributeProto.STRING: lambda a: a.s.decode("utf-8", errors="replace"),
+    AttributeProto.FLOATS: lambda a: list(a.floats),
+    AttributeProto.INTS: lambda a: list(a.ints),
+    AttributeProto.STRINGS: lambda a: [
+        s.decode("utf-8", errors="replace") for s in a.strings
+    ],
+    AttributeProto.TENSOR: lambda a: numpy_helper.to_array(a.t).tolist(),
+    AttributeProto.TENSORS: lambda a: [
+        numpy_helper.to_array(t).tolist() for t in a.tensors
+    ],
+}
+
+
+def parse_attribute(
+    attr: AttributeProto,
+) -> float | int | str | list[int] | list[float] | list[str]:
+    """Parse ONNX attributes into a Python-native type.
+
+    Args:
+        attr (AttributeProto): The ONNX attribute to parse.
+
+    Raises:
+        ValueError: If the attribute type is unsupported.
+
+    Returns:
+        Any: The parsed attribute value as a Python type.
+    """
+    parser = ATTRIBUTE_PARSERS.get(attr.type)
+    if parser is None:
+        msg = f"Unsupported attribute type: {attr.type}"
+        raise ValueError(msg)
+    return parser(attr)
+
+
+def parse_attributes(attrs: list[AttributeProto]) -> dict[str, Any]:
+    """Parse multiple ONNX attributes into a dictionary.
+
+    Args:
+        attrs (list[AttributeProto]): List of ONNX attributes.
+
+    Returns:
+        dict[str, Any]: Mapping of attribute names to their parsed values.
+    """
+    return {attr.name: parse_attribute(attr) for attr in attrs}
+
+
+def extract_shape_dict(inferred_model: onnx.GraphProto) -> dict[str, list[int]]:
+    """Extract shape information from an ONNX model's graph.
+
+    Args:
+        inferred_model (onnx.GraphProto): The inferred ONNX model graph.
+
+    Returns:
+        dict[str, list[int]]: Mapping from tensor names to their shape dimensions.
+                              Unknown dimensions are returned as 1.
+    """
+    value_info = {}
+    graph = inferred_model.graph
+    all_info = list(graph.value_info) + list(graph.output) + list(graph.input)
+    for vi in all_info:
+        if vi.type.HasField("tensor_type"):
+            shape = [
+                # TODO@jsgold-1: figure out how to deal with bad value # noqa: FIX002, TD003, E501
+                d.dim_value if d.HasField("dim_value") else 1
+                for d in vi.type.tensor_type.shape.dim
+            ]
+            value_info[vi.name] = shape
+    return value_info
+
+
+def replace_input_references(
+    graph: onnx.GraphProto,
+    old_output: str,
+    new_output: str,
+) -> None:
+    """Replace all references to an input tensor in an ONNX graph.
+
+    Args:
+        graph (onnx.GraphProto): The ONNX graph to modify.
+        old_output (str): The original tensor name to replace.
+        new_output (str): The new tensor name.
+    """
+    for node in graph.node:
+        for i, input_name in enumerate(node.input):
+            if input_name == old_output:
+                node.input[i] = new_output
+
+
+def extract_attributes(node: onnx.NodeProto) -> dict:
+    """Extract all attributes from an ONNX node into a Python dictionary.
+
+    Args:
+        node (onnx.NodeProto): The ONNX node to extract attributes from.
+
+    Raises:
+        ValueError: If an attribute type is unsupported.
+
+    Returns:
+        dict: Mapping of attribute names to Python-native values.
+    """
+    attrs = {}
+    for attr in node.attribute:
+        name = attr.name
+        val = onnx.helper.get_attribute_value(attr)
+
+        if attr.type == AttributeProto.FLOAT:
+            attrs[name] = float(val)
+        elif attr.type == AttributeProto.INT:
+            attrs[name] = int(val)
+        elif attr.type == AttributeProto.FLOATS:
+            attrs[name] = [
+                float(x) for x in val
+            ]  # ← you want to ensure these are int if your op expects it
+        elif attr.type == AttributeProto.INTS:
+            attrs[name] = ",".join(str(v) for v in val)
+        elif attr.type == AttributeProto.STRING:
+            attrs[name] = val.decode("utf-8") if isinstance(val, bytes) else val
+        elif attr.type == AttributeProto.BOOL:
+            attrs[name] = bool(val)
+        else:
+            msg = f"Unsupported attribute type: {attr.name} (type={attr.type})"
+            raise ValueError(msg)
+    return attrs
+
+
+def get_input_shapes(onnx_model: onnx.ModelProto) -> dict:
+    """Get the input tensor shapes from an ONNX model.
+
+    Args:
+        onnx_model (onnx.ModelProto): The ONNX model.
+
+    Returns:
+        dict: Mapping from input tensor names to their shape dimensions.
+    """
+    input_shapes = {}
+    for model_in in onnx_model.graph.input:
+        input_name = model_in.name
+        # Get the shape from the input's type information
+        shape = [dim.dim_value for dim in model_in.type.tensor_type.shape.dim]
+        input_shapes[input_name] = shape
+    return input_shapes
+
+
+def get_attribute_ints(
+    node: onnx.NodeProto,
+    name: str,
+    default: list[int] | None = None,
+) -> list[int]:
+    """Retrieve a list of integer values from an ONNX node's attribute.
+
+    Args:
+        node (onnx.NodeProto): The ONNX node.
+        name (str): Name of the attribute to retrieve.
+        default (list[int], optional):
+            Default list to return if the attribute is missing. Defaults to None.
+
+    Returns:
+        list[int]: List of integers from the attribute, or the default if not found.
+    """
+    for attr in node.attribute:
+        if attr.name == name and attr.type == onnx.AttributeProto.INTS:
+            return list(attr.ints)
+    return default if default is not None else []