onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/op_registry.py

@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ONNX operator registry with custom operator and opset version support."""
+
+import copy
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Dict, Callable, Optional, Union, List, Tuple
+
+import onnx
+
+if TYPE_CHECKING:  # pragma: no cover - only for type checking
+    from .graph_builder import GraphBuilder
+
+import torch.fx
+
+# Type alias for operator handler functions.
+# Handlers take a GraphBuilder and ONNX node, returning one or more FX nodes.
+OpHandler = Callable[
+    ["GraphBuilder", onnx.NodeProto], Union[torch.fx.Node, Tuple[torch.fx.Node, ...]]
+]
+
+# Registry: {domain: {op_type: [(since_version, handler), ...]}}
+# Handlers are stored in descending version order for efficient lookup.
+# Empty string "" represents the default ONNX domain.
+_VERSIONED_REGISTRY: Dict[str, Dict[str, List[Tuple[int, OpHandler]]]] = {"": {}}
+
+
+@contextmanager
+def registry_context():
+    """Temporarily isolate registry mutations (intended for tests)."""
+    snapshot = copy.deepcopy(_VERSIONED_REGISTRY)
+    try:
+        yield
+    finally:
+        _VERSIONED_REGISTRY.clear()
+        _VERSIONED_REGISTRY.update(snapshot)
+
+
+def register(
+    op_type: str, domain: str = "", since_version: int = 1
+) -> Callable[[OpHandler], OpHandler]:
+    """Decorator to register an ONNX operator handler with version support.
+
+    Parameters
+    ----------
+    op_type : str
+        The ONNX operator type name (e.g., "Add", "Relu").
+    domain : str, optional
+        The ONNX domain (e.g., "com.microsoft"). Default is "" (standard ONNX domain).
+    since_version : int, optional
+        The minimum opset version this handler supports. Default is 1.
+
+    Returns
+    -------
+    Callable
+        Decorator function.
+
+    Examples
+    --------
+    Register a standard ONNX operator (all versions):
+
+    >>> @register("MyOp")
+    ... def my_op(builder, node):
+    ...     x = builder.get_value(node.input[0])
+    ...     return builder.call_function(torch.relu, args=(x,))
+
+    Register version-specific handlers:
+
+    >>> @register("Softmax", since_version=1)
+    ... def softmax_v1(builder, node):
+    ...     # opset 1-12: axis defaults to 1
+    ...     ...
+
+    >>> @register("Softmax", since_version=13)
+    ... def softmax_v13(builder, node):
+    ...     # opset 13+: axis defaults to -1
+    ...     ...
+
+    Register a custom domain operator:
+
+    >>> @register("CustomOp", domain="com.mycompany")
+    ... def custom_op(builder, node):
+    ...     x = builder.get_value(node.input[0])
+    ...     return builder.call_function(my_custom_function, args=(x,))
+    """
+
+    def decorator(func: OpHandler) -> OpHandler:
+        if domain not in _VERSIONED_REGISTRY:
+            _VERSIONED_REGISTRY[domain] = {}
+        if op_type not in _VERSIONED_REGISTRY[domain]:
+            _VERSIONED_REGISTRY[domain][op_type] = []
+
+        handlers = _VERSIONED_REGISTRY[domain][op_type]
+        # Remove existing handler with same since_version to allow re-registration
+        handlers[:] = [(v, h) for v, h in handlers if v != since_version]
+        handlers.append((since_version, func))
+        # Keep sorted in descending order by version for efficient lookup
+        handlers.sort(key=lambda x: x[0], reverse=True)
+
+        return func
+
+    return decorator
+
+
+def register_op(
+    op_type: str,
+    handler: Optional[OpHandler] = None,
+    domain: str = "",
+    since_version: int = 1,
+) -> Union[OpHandler, Callable[[OpHandler], OpHandler]]:
+    """Register an ONNX operator handler with version support.
+
+    This function can be used as a decorator or called directly to register
+    custom operator handlers for ONNX operators that are not natively supported.
+
+    Parameters
+    ----------
+    op_type : str
+        The ONNX operator type name.
+    handler : OpHandler, optional
+        The handler function. If not provided, returns a decorator.
+    domain : str, optional
+        The ONNX domain. Default is "" (standard ONNX domain).
+    since_version : int, optional
+        The minimum opset version this handler supports. Default is 1.
+
+    Returns
+    -------
+    OpHandler or Callable
+        If handler is provided, returns the handler.
+        Otherwise, returns a decorator function.
+
+    Examples
+    --------
+    Using as a decorator:
+
+    >>> @register_op("MyCustomOp")
+    ... def my_custom_op(builder, node):
+    ...     x = builder.get_value(node.input[0])
+    ...     return builder.call_function(torch.sigmoid, args=(x,))
+
+    Using as a function:
+
+    >>> def my_handler(builder, node):
+    ...     x = builder.get_value(node.input[0])
+    ...     return builder.call_function(torch.tanh, args=(x,))
+    >>> register_op("TanhCustom", my_handler)
+
+    Registering for a custom domain:
+
+    >>> @register_op("BiasGelu", domain="com.microsoft")
+    ... def bias_gelu(builder, node):
+    ...     x = builder.get_value(node.input[0])
+    ...     bias = builder.get_value(node.input[1])
+    ...     return builder.call_function(
+    ...         lambda t, b: torch.nn.functional.gelu(t + b),
+    ...         args=(x, bias)
+    ...     )
+
+    Registering version-specific handlers:
+
+    >>> @register_op("MyOp", since_version=1)
+    ... def my_op_v1(builder, node): ...
+
+    >>> @register_op("MyOp", since_version=13)
+    ... def my_op_v13(builder, node): ...
+    """
+    if handler is not None:
+        # Direct call: register_op("Op", handler)
+        if domain not in _VERSIONED_REGISTRY:
+            _VERSIONED_REGISTRY[domain] = {}
+        if op_type not in _VERSIONED_REGISTRY[domain]:
+            _VERSIONED_REGISTRY[domain][op_type] = []
+
+        handlers = _VERSIONED_REGISTRY[domain][op_type]
+        # Remove existing handler with same since_version
+        handlers[:] = [(v, h) for v, h in handlers if v != since_version]
+        handlers.append((since_version, handler))
+        handlers.sort(key=lambda x: x[0], reverse=True)
+
+        return handler
+    else:
+        # Decorator usage: @register_op("Op")
+        return register(op_type, domain, since_version)
+
+
+def unregister_op(
+    op_type: str, domain: str = "", since_version: Optional[int] = None
+) -> bool:
+    """Unregister an operator handler.
+
+    Parameters
+    ----------
+    op_type : str
+        The ONNX operator type name.
+    domain : str, optional
+        The ONNX domain. Default is "" (standard ONNX domain).
+    since_version : int, optional
+        The specific version handler to remove. If None, removes all versions.
+
+    Returns
+    -------
+    bool
+        True if the operator was unregistered, False if it wasn't registered.
+    """
+    if domain not in _VERSIONED_REGISTRY:
+        return False
+    if op_type not in _VERSIONED_REGISTRY[domain]:
+        return False
+
+    handlers = _VERSIONED_REGISTRY[domain][op_type]
+    if since_version is None:
+        # Remove all versions
+        del _VERSIONED_REGISTRY[domain][op_type]
+        return True
+    else:
+        # Remove specific version
+        original_len = len(handlers)
+        handlers[:] = [(v, h) for v, h in handlers if v != since_version]
+        if len(handlers) < original_len:
+            if not handlers:
+                del _VERSIONED_REGISTRY[domain][op_type]
+            return True
+        return False
+
+
+def get_handler(
+    op_type: str, domain: str = "", opset_version: int = 23
+) -> Optional[OpHandler]:
+    """Get the handler for an operator at a specific opset version.
+
+    Finds the handler with the highest since_version that is <= opset_version.
+
+    Parameters
+    ----------
+    op_type : str
+        The ONNX operator type name.
+    domain : str, optional
+        The ONNX domain. Default is "" (standard ONNX domain).
+    opset_version : int, optional
+        The target opset version. Default is 23 (current latest).
+
+    Returns
+    -------
+    OpHandler or None
+        The appropriate handler function, or None if not found.
+    """
+    # Normalize domain: "ai.onnx" is equivalent to ""
+    if domain == "ai.onnx":
+        domain = ""
+
+    if domain not in _VERSIONED_REGISTRY:
+        return None
+
+    handlers = _VERSIONED_REGISTRY[domain].get(op_type)
+    if not handlers:
+        return None
+
+    # Handlers are sorted in descending order by since_version
+    # Find the first handler where since_version <= opset_version
+    for since_version, handler in handlers:
+        if since_version <= opset_version:
+            return handler
+
+    return None
+
+
+def is_supported(op_type: str, domain: str = "", opset_version: int = 23) -> bool:
+    """Check if an operator is supported.
+
+    Parameters
+    ----------
+    op_type : str
+        The ONNX operator type name.
+    domain : str, optional
+        The ONNX domain. Default is "" (standard ONNX domain).
+    opset_version : int, optional
+        The target opset version. Default is 23.
+
+    Returns
+    -------
+    bool
+        True if the operator is supported.
+    """
+    return get_handler(op_type, domain, opset_version) is not None
+
+
+def get_supported_ops(domain: str = "") -> list:
+    """Get list of supported ONNX operators for a domain.
+
+    Parameters
+    ----------
+    domain : str, optional
+        The ONNX domain. Default is "" (standard ONNX domain).
+
+    Returns
+    -------
+    list
+        Sorted list of supported operator names.
+    """
+    if domain in _VERSIONED_REGISTRY:
+        return sorted(_VERSIONED_REGISTRY[domain].keys())
+    return []
+
+
+def get_all_supported_ops() -> Dict[str, list]:
+    """Get all supported operators across all domains.
+
+    Returns
+    -------
+    Dict[str, list]
+        Dictionary mapping domain names to sorted lists of operator names.
+    """
+    return {domain: sorted(ops.keys()) for domain, ops in _VERSIONED_REGISTRY.items()}
+
+
+def get_registered_domains() -> list:
+    """Get list of registered domains.
+
+    Returns
+    -------
+    list
+        List of domain names.
+    """
+    return list(_VERSIONED_REGISTRY.keys())
+
+
+def get_handler_versions(op_type: str, domain: str = "") -> List[int]:
+    """Get all registered opset versions for an operator.
+
+    Parameters
+    ----------
+    op_type : str
+        The ONNX operator type name.
+    domain : str, optional
+        The ONNX domain. Default is "" (standard ONNX domain).
+
+    Returns
+    -------
+    List[int]
+        List of registered since_version values, sorted in ascending order.
+    """
+    if domain in _VERSIONED_REGISTRY:
+        handlers = _VERSIONED_REGISTRY[domain].get(op_type, [])
+        return sorted([v for v, _ in handlers])
+    return []

onnx2fx/ops/__init__.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ONNX operator implementations.
+
+This package contains all ONNX operator implementations organized by category:
+
+- activation.py: Activation functions (Relu, Sigmoid, Softmax, etc.)
+- arithmetic.py: Arithmetic and math ops (Add, Mul, Sin, Cos, etc.)
+- attention.py: Attention mechanisms (standard ONNX domain)
+- attention_msft.py: Attention mechanisms (com.microsoft domain)
+- control_flow.py: Control flow ops (Loop, If, Scan)
+- convolution.py: Convolution ops (Conv, ConvTranspose, DeformConv)
+- image.py: Image processing ops (Resize, DepthToSpace, etc.)
+- linalg.py: Linear algebra ops (Einsum, Det)
+- loss.py: Loss functions (SoftmaxCrossEntropyLoss, etc.)
+- nn.py: Core neural network ops (MatMul, Gemm, Dropout)
+- normalization.py: Normalization ops (BatchNorm, LayerNorm, etc.)
+- pooling.py: Pooling ops (MaxPool, AveragePool, etc.)
+- quantization.py: Quantization ops (QLinearConv, etc.)
+- random.py: Random number generation (RandomNormal, etc.)
+- recurrent.py: Recurrent neural networks (LSTM, GRU, RNN)
+- reduction.py: Reduction ops (ReduceSum, ReduceMean, etc.)
+- sequence.py: Sequence ops (SequenceConstruct, etc.)
+- signal.py: Signal processing (STFT, MelWeightMatrix, window functions, NMS)
+- string.py: String ops (StringNormalizer)
+- tensor.py: Tensor manipulation ops (Reshape, Transpose, etc.)
+- training.py: Training ops (Gradient, Momentum, Adagrad)
+"""
+
+# Import all operator modules to register handlers
+from . import activation
+from . import arithmetic
+from . import attention
+from . import attention_msft
+from . import control_flow
+from . import convolution
+from . import image
+from . import linalg
+from . import loss
+from . import nn
+from . import normalization
+from . import pooling
+from . import quantization
+from . import random
+from . import recurrent
+from . import reduction
+from . import sequence
+from . import signal
+from . import string
+from . import tensor
+from . import training
+
+__all__ = [
+    "activation",
+    "arithmetic",
+    "attention",
+    "attention_msft",
+    "control_flow",
+    "convolution",
+    "image",
+    "linalg",
+    "loss",
+    "nn",
+    "normalization",
+    "pooling",
+    "quantization",
+    "random",
+    "recurrent",
+    "reduction",
+    "sequence",
+    "signal",
+    "string",
+    "tensor",
+    "training",
+]

onnx2fx/ops/activation.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Activation function operators."""
+
+from typing import TYPE_CHECKING, Callable
+
+import onnx
+import torch
+import torch.nn.functional as F
+
+from ..op_registry import register
+from ..utils.attributes import get_attribute
+from ..utils.op_helpers import unary_op, unary_op_with_kwargs
+
+if TYPE_CHECKING:
+    from ..graph_builder import GraphBuilder
+
+
+def _coerced_softmax_handler(
+    torch_fn: Callable[..., torch.Tensor],
+    *,
+    default_axis: int,
+    doc: str,
+) -> Callable[["GraphBuilder", onnx.NodeProto], torch.fx.Node]:
+    def handler(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+        x = builder.get_value(node.input[0])
+        axis = get_attribute(node, "axis", default_axis)
+
+        def _coerced_softmax(t: torch.Tensor, axis: int) -> torch.Tensor:
+            # Handle negative axis
+            if axis < 0:
+                axis = t.dim() + axis
+
+            # Coerce to 2D: flatten [0:axis] and [axis:]
+            orig_shape = t.shape
+            pre_dim = 1
+            for i in range(axis):
+                pre_dim *= t.shape[i]
+
+            t_2d = t.reshape(pre_dim, -1)
+            result_2d = torch_fn(t_2d, dim=1)
+            return result_2d.reshape(orig_shape)
+
+        return builder.call_function(_coerced_softmax, args=(x, axis))
+
+    handler.__doc__ = doc
+    return handler
+
+
+register("Relu")(unary_op(F.relu, "ReLU activation."))
+
+
+register("LeakyRelu")(
+    unary_op_with_kwargs(
+        F.leaky_relu,
+        attr_map={"negative_slope": ("alpha", 0.01)},
+        doc="Leaky ReLU activation.",
+    )
+)
+
+
+@register("PRelu")
+def prelu(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Parametric ReLU activation.
+
+    ONNX PRelu allows slope to be broadcastable to the input tensor,
+    while PyTorch's F.prelu expects slope to match channel dimension.
+    We implement using torch.where for proper broadcasting support.
+
+    When slope has shape [C] and input has shape [N, C, ...], we need to
+    reshape slope to [1, C, 1, ...] for proper broadcasting.
+    """
+    x = builder.get_value(node.input[0])
+    slope = builder.get_value(node.input[1])
+
+    def _prelu(x: torch.Tensor, slope: torch.Tensor) -> torch.Tensor:
+        # If slope is 1D with size matching channels and input is ND with N > 1,
+        # reshape slope for proper broadcasting along channel dimension (dim=1)
+        if slope.ndim == 1 and x.ndim > 1 and slope.numel() == x.shape[1]:
+            # Reshape [C] to [1, C, 1, 1, ...] for broadcasting
+            shape = [1, slope.numel()] + [1] * (x.ndim - 2)
+            slope = slope.view(shape)
+        return torch.where(x >= 0, x, x * slope)
+
+    return builder.call_function(_prelu, args=(x, slope))
+
+
+register("Elu")(
+    unary_op_with_kwargs(
+        F.elu,
+        attr_map={"alpha": ("alpha", 1.0)},
+        doc="Exponential Linear Unit activation.",
+    )
+)
+
+
+@register("Selu")
+def selu(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Scaled Exponential Linear Unit activation.
+
+    ONNX SELU: y = gamma * (alpha * (exp(x) - 1) for x < 0, x for x >= 0)
+    PyTorch SELU uses fixed alpha=1.6732... and gamma=1.0507...
+    ONNX defaults: alpha=1.67326..., gamma=1.0507... but allows custom values.
+    """
+    x = builder.get_value(node.input[0])
+    alpha = get_attribute(node, "alpha", 1.67326319217681884765625)
+    gamma = get_attribute(node, "gamma", 1.05070102214813232421875)
+
+    # PyTorch's fixed SELU values
+    pytorch_alpha = 1.6732632423543772848170429916717
+    pytorch_gamma = 1.0507009873554804934193349852946
+
+    # If using PyTorch's fixed values (within tolerance), use F.selu for efficiency
+    if abs(alpha - pytorch_alpha) < 1e-5 and abs(gamma - pytorch_gamma) < 1e-5:
+        return builder.call_function(F.selu, args=(x,))
+
+    # Otherwise implement manually: gamma * (alpha * (exp(x) - 1) for x < 0, x for x >= 0)
+    def _custom_selu(x: torch.Tensor, alpha: float, gamma: float) -> torch.Tensor:
+        return gamma * torch.where(x > 0, x, alpha * (torch.exp(x) - 1))
+
+    return builder.call_function(_custom_selu, args=(x, alpha, gamma))
+
+
+register("Celu")(
+    unary_op_with_kwargs(
+        F.celu,
+        attr_map={"alpha": ("alpha", 1.0)},
+        doc="Continuously Differentiable Exponential Linear Unit activation.",
+    )
+)
+
+
+register("Sigmoid")(unary_op(torch.sigmoid, "Sigmoid activation."))
+
+
+@register("HardSigmoid")
+def hard_sigmoid(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Hard Sigmoid activation.
+
+    ONNX HardSigmoid: max(0, min(1, alpha * x + beta))
+    ONNX defaults: alpha=0.2, beta=0.5
+    PyTorch hardsigmoid uses fixed alpha=1/6, beta=0.5
+    """
+    x = builder.get_value(node.input[0])
+    alpha = get_attribute(node, "alpha", 0.2)
+    beta = get_attribute(node, "beta", 0.5)
+
+    # PyTorch hardsigmoid uses alpha=1/6 ≈ 0.16667, beta=0.5
+    pytorch_alpha = 1.0 / 6.0
+
+    # If using PyTorch's fixed values (within tolerance), use F.hardsigmoid
+    if abs(alpha - pytorch_alpha) < 1e-5 and abs(beta - 0.5) < 1e-5:
+        return builder.call_function(F.hardsigmoid, args=(x,))
+
+    # Otherwise implement manually: max(0, min(1, alpha * x + beta))
+    def _custom_hardsigmoid(x: torch.Tensor, alpha: float, beta: float) -> torch.Tensor:
+        return torch.clamp(alpha * x + beta, 0.0, 1.0)
+
+    return builder.call_function(_custom_hardsigmoid, args=(x, alpha, beta))
+
+
+register("Tanh")(unary_op(torch.tanh, "Tanh activation."))
+
+
+register("Softmax", since_version=1)(
+    _coerced_softmax_handler(
+        F.softmax,
+        default_axis=1,
+        doc=(
+            "Softmax activation for opset 1-12 with axis defaulting to 1 and "
+            "2D coercion."
+        ),
+    )
+)
+
+
+register("Softmax", since_version=13)(
+    unary_op_with_kwargs(
+        F.softmax,
+        attr_map={"dim": ("axis", -1)},
+        doc=(
+            "Softmax activation for opset 13+ with axis defaulting to the last "
+            "dimension."
+        ),
+    )
+)
+
+
+register("LogSoftmax", since_version=1)(
+    _coerced_softmax_handler(
+        F.log_softmax,
+        default_axis=1,
+        doc=(
+            "Log Softmax activation for opset 1-12 with axis defaulting to 1 "
+            "and 2D coercion."
+        ),
+    )
+)
+
+
+register("LogSoftmax", since_version=13)(
+    unary_op_with_kwargs(
+        F.log_softmax,
+        attr_map={"dim": ("axis", -1)},
+        doc="Log Softmax activation for opset 13+.",
+    )
+)
+
+
+register("Softplus")(unary_op(F.softplus, "Softplus activation."))
+register("Softsign")(unary_op(F.softsign, "Softsign activation."))
+
+
+register("Gelu")(
+    unary_op_with_kwargs(
+        F.gelu,
+        attr_map={"approximate": ("approximate", "none")},
+        doc="Gaussian Error Linear Unit activation.",
+    )
+)
+
+
+register("Silu")(unary_op(F.silu, "Sigmoid Linear Unit (SiLU/Swish) activation."))
+register("Swish")(unary_op(F.silu, "Swish activation (alias for SiLU)."))
+register("Mish")(unary_op(F.mish, "Mish activation."))
+
+
+register("ThresholdedRelu")(
+    unary_op_with_kwargs(
+        F.threshold,
+        attr_map={"threshold": ("alpha", 1.0)},
+        fixed_kwargs={"value": 0.0},
+        doc="Thresholded ReLU activation.",
+    )
+)
+
+
+register("HardSwish")(unary_op(F.hardswish, "Hard Swish activation."))
+
+
+@register("Hardmax")
+def hardmax(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Hardmax - one-hot encoding of argmax."""
+    x = builder.get_value(node.input[0])
+    axis = get_attribute(node, "axis", -1)
+
+    def _hardmax(t: torch.Tensor, ax: int) -> torch.Tensor:
+        # Normalize axis to positive
+        if ax < 0:
+            ax = t.dim() + ax
+        # one_hot appends the class dimension at the end
+        one_hot = torch.nn.functional.one_hot(
+            torch.argmax(t, dim=ax), num_classes=t.shape[ax]
+        ).to(t.dtype)
+        # Move the one-hot dimension from the end back to the original axis position
+        # one_hot has shape: [...dims before ax..., ...dims after ax..., num_classes]
+        # We need to move the last dim to position ax
+        return torch.movedim(one_hot, -1, ax)
+
+    return builder.call_function(_hardmax, args=(x, axis))
+
+
+@register("Shrink")
+def shrink(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Shrink activation.
+
+    If x < -lambd: y = x + bias
+    If x > lambd: y = x - bias
+    Otherwise: y = 0
+    """
+    x = builder.get_value(node.input[0])
+    bias = get_attribute(node, "bias", 0.0)
+    lambd = get_attribute(node, "lambd", 0.5)
+
+    def _shrink(t: torch.Tensor, bias: float, lambd: float) -> torch.Tensor:
+        result = torch.zeros_like(t)
+        mask_neg = t < -lambd
+        mask_pos = t > lambd
+        result = torch.where(mask_neg, t + bias, result)
+        result = torch.where(mask_pos, t - bias, result)
+        return result
+
+    return builder.call_function(_shrink, args=(x, bias, lambd))