emx-onnx-cgen 0.3.7__py3-none-any.whl → 0.4.1.dev0__py3-none-any.whl

emx_onnx_cgen/lowering/dequantize_linear.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import normalize_axis
+from .common import (
+    optional_name,
+    value_dtype as _value_dtype,
+    value_shape as _value_shape,
+)
+from .registry import register_lowering
+from ..ir.ops import DequantizeLinearOp
+
+
+@dataclass(frozen=True)
+class DequantizeSpec:
+    input_shape: tuple[int, ...]
+    scale_shape: tuple[int, ...]
+    axis: int | None
+    block_size: int | None
+
+
+def resolve_dequantize_spec(graph: Graph, node: Node) -> DequantizeSpec:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "DequantizeLinear must have 2 or 3 inputs and 1 output"
+        )
+    supported_attrs = {"axis", "block_size"}
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("DequantizeLinear has unsupported attributes")
+    block_size = int(node.attrs.get("block_size", 0))
+    if block_size < 0:
+        raise UnsupportedOpError("DequantizeLinear block_size must be >= 0")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    scale_shape = _value_shape(graph, node.inputs[1], node)
+    zero_point_name = optional_name(node.inputs, 2)
+    if zero_point_name is not None:
+        zero_point_shape = _value_shape(graph, zero_point_name, node)
+        if zero_point_shape != scale_shape:
+            raise ShapeInferenceError(
+                "DequantizeLinear zero_point shape must match scale shape"
+            )
+    if scale_shape not in {(), (1,)}:
+        axis = int(node.attrs.get("axis", 1))
+        axis = normalize_axis(axis, input_shape, node)
+        if block_size > 0:
+            if len(scale_shape) != len(input_shape):
+                raise UnsupportedOpError(
+                    "DequantizeLinear blocked scales must match input rank"
+                )
+            if input_shape[axis] % block_size != 0:
+                raise ShapeInferenceError(
+                    "DequantizeLinear block_size must evenly divide axis length"
+                )
+            expected = list(input_shape)
+            expected[axis] = input_shape[axis] // block_size
+            if scale_shape != tuple(expected):
+                raise ShapeInferenceError(
+                    "DequantizeLinear blocked scale shape must match "
+                    "input shape with a reduced axis"
+                )
+        else:
+            if len(scale_shape) != 1:
+                raise UnsupportedOpError(
+                    "DequantizeLinear supports per-tensor, per-axis, "
+                    "and blocked scales only"
+                )
+            if scale_shape[0] != input_shape[axis]:
+                raise ShapeInferenceError(
+                    "DequantizeLinear scale length must match input axis size"
+                )
+    else:
+        axis = None
+        block_size = 0
+    return DequantizeSpec(
+        input_shape=input_shape,
+        scale_shape=scale_shape,
+        axis=axis,
+        block_size=block_size or None,
+    )
+
+
+@register_lowering("DequantizeLinear")
+def lower_dequantize_linear(graph: Graph, node: Node) -> DequantizeLinearOp:
+    input_dtype = _value_dtype(graph, node.inputs[0], node)
+    scale_dtype = _value_dtype(graph, node.inputs[1], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if input_dtype not in {
+        ScalarType.U8,
+        ScalarType.I8,
+        ScalarType.U16,
+        ScalarType.I16,
+    }:
+        raise UnsupportedOpError(
+            "DequantizeLinear supports int8/uint8/int16/uint16 inputs only"
+        )
+    if not scale_dtype.is_float or not output_dtype.is_float:
+        raise UnsupportedOpError(
+            "DequantizeLinear supports float16/float/double scales and outputs only"
+        )
+    if output_dtype != scale_dtype:
+        raise UnsupportedOpError(
+            "DequantizeLinear output dtype must match scale dtype"
+        )
+    zero_point_name = optional_name(node.inputs, 2)
+    if zero_point_name is not None:
+        zero_point_dtype = _value_dtype(graph, zero_point_name, node)
+        if zero_point_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "DequantizeLinear zero_point dtype must match input dtype"
+            )
+    spec = resolve_dequantize_spec(graph, node)
+    return DequantizeLinearOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        zero_point=zero_point_name,
+        output=node.outputs[0],
+        input_shape=spec.input_shape,
+        axis=spec.axis,
+        block_size=spec.block_size,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+        scale_dtype=scale_dtype,
+    )

emx_onnx_cgen/lowering/elementwise.py

@@ -3,11 +3,18 @@ from __future__ import annotations
 from shared.scalar_functions import ScalarFunction, ScalarFunctionError
 from shared.scalar_types import ScalarType
 
-from ..ir.ops import BinaryOp, ClipOp, UnaryOp
+from ..ir.op_base import BroadcastingOpBase
+from ..ir.ops import BinaryOp, ClipOp, PowOp, UnaryOp
 from ..errors import UnsupportedOpError
 from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
-from ..lowering.common import node_dtype, optional_name, value_dtype, value_shape
+from ..lowering.common import (
+    node_dtype,
+    onnx_opset_version,
+    optional_name,
+    value_dtype,
+    value_shape,
+)
 from ..lowering.registry import register_lowering, register_lowering_if_missing
 from ..ops import (
     BINARY_OP_TYPES,
@@ -29,6 +36,24 @@ def lower_clip(graph: Graph, node: Node) -> ClipOp:
         raise UnsupportedOpError("Clip input must be provided")
     min_name = optional_name(node.inputs, 1)
     max_name = optional_name(node.inputs, 2)
+    min_value = None
+    max_value = None
+    opset_version = onnx_opset_version(graph)
+    if opset_version is None or opset_version < 11:
+        if min_name is None and "min" in node.attrs:
+            try:
+                min_value = float(node.attrs["min"])
+            except (TypeError, ValueError) as exc:
+                raise UnsupportedOpError(
+                    "Clip min attribute must be numeric"
+                ) from exc
+        if max_name is None and "max" in node.attrs:
+            try:
+                max_value = float(node.attrs["max"])
+            except (TypeError, ValueError) as exc:
+                raise UnsupportedOpError(
+                    "Clip max attribute must be numeric"
+                ) from exc
     input_dtype = value_dtype(graph, input_name, node)
     output_dtype = value_dtype(graph, node.outputs[0], node)
     if input_dtype != output_dtype:
@@ -61,11 +86,8 @@ def lower_clip(graph: Graph, node: Node) -> ClipOp:
         input_min=min_name,
         input_max=max_name,
         output=node.outputs[0],
-        input_shape=input_shape,
-        min_shape=min_shape,
-        max_shape=max_shape,
-        output_shape=output_shape,
-        dtype=input_dtype,
+        min_value=min_value,
+        max_value=max_value,
     )
 
 
@@ -82,9 +104,54 @@ def lower_celu(graph: Graph, node: Node) -> UnaryOp:
         input0=node.inputs[0],
         output=node.outputs[0],
         function=ScalarFunction.CELU,
-        shape=output_shape,
-        dtype=dtype,
-        input_dtype=dtype,
+        params=(alpha,),
+    )
+
+
+@register_lowering("Elu")
+def lower_elu(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Elu must have 1 input and 1 output")
+    dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not dtype.is_float:
+        raise UnsupportedOpError("Elu only supports floating-point inputs")
+    for key in node.attrs:
+        if key != "alpha":
+            raise UnsupportedOpError(f"Elu does not support attribute {key}")
+    try:
+        alpha = float(node.attrs.get("alpha", 1.0))
+    except (TypeError, ValueError) as exc:
+        raise UnsupportedOpError("Elu alpha must be numeric") from exc
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.ELU,
+        params=(alpha,),
+    )
+
+
+@register_lowering("LeakyRelu")
+def lower_leaky_relu(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("LeakyRelu must have 1 input and 1 output")
+    dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not dtype.is_float:
+        raise UnsupportedOpError("LeakyRelu only supports floating-point inputs")
+    for key in node.attrs:
+        if key != "alpha":
+            raise UnsupportedOpError(
+                f"LeakyRelu does not support attribute {key}"
+            )
+    try:
+        alpha = float(node.attrs.get("alpha", 0.01))
+    except (TypeError, ValueError) as exc:
+        raise UnsupportedOpError("LeakyRelu alpha must be numeric") from exc
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.LEAKY_RELU,
         params=(alpha,),
     )
 
@@ -102,9 +169,6 @@ def lower_swish(graph: Graph, node: Node) -> UnaryOp:
         input0=node.inputs[0],
         output=node.outputs[0],
         function=ScalarFunction.SWISH,
-        shape=output_shape,
-        dtype=dtype,
-        input_dtype=dtype,
         params=(alpha,),
     )
 
@@ -123,13 +187,50 @@ def lower_shrink(graph: Graph, node: Node) -> UnaryOp:
         input0=node.inputs[0],
         output=node.outputs[0],
         function=ScalarFunction.SHRINK,
-        shape=output_shape,
-        dtype=dtype,
-        input_dtype=dtype,
         params=(bias, lambd),
     )
 
 
+@register_lowering("Pow")
+def lower_pow(graph: Graph, node: Node) -> PowOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Pow must have 2 inputs and 1 output")
+    op_dtype = value_dtype(graph, node.inputs[0], node)
+    op_spec = binary_op_symbol(ScalarFunction.POW, node.attrs, dtype=op_dtype)
+    if op_spec is None:
+        raise UnsupportedOpError("Unsupported op Pow")
+    return PowOp(
+        input0=node.inputs[0],
+        input1=node.inputs[1],
+        output=node.outputs[0],
+        function=ScalarFunction.POW,
+        operator_kind=op_spec.kind,
+    )
+
+
+def _infer_binary_output_shape(
+    *,
+    function: ScalarFunction,
+    input0_shape: tuple[int, ...],
+    input1_shape: tuple[int, ...],
+) -> tuple[int, ...]:
+    if function != ScalarFunction.PRELU:
+        return BroadcastingOpBase.broadcast_shapes(input0_shape, input1_shape)
+    if BroadcastingOpBase.unidirectional_broadcastable(
+        input1_shape, input0_shape
+    ):
+        return input0_shape
+    channel_axis = BroadcastingOpBase.prelu_channel_axis(
+        input0_shape, input1_shape
+    )
+    if channel_axis is None:
+        raise ShapeInferenceError(
+            "Broadcasting mismatch for shapes: "
+            + ", ".join(str(shape) for shape in (input0_shape, input1_shape))
+        )
+    return input0_shape
+
+
 def _lower_binary_unary(graph: Graph | GraphContext, node: Node) -> BinaryOp | UnaryOp:
     if node.op_type == "BitShift":
         if len(node.inputs) != 2 or len(node.outputs) != 1:
@@ -163,11 +264,6 @@ def _lower_binary_unary(graph: Graph | GraphContext, node: Node) -> BinaryOp | U
             output=node.outputs[0],
             function=function,
             operator_kind=op_spec.kind,
-            input0_shape=input0_shape,
-            input1_shape=input1_shape,
-            shape=output_shape,
-            dtype=op_dtype,
-            input_dtype=op_dtype,
         )
     if node.op_type == "Mod":
         fmod = int(node.attrs.get("fmod", 0))
@@ -201,18 +297,21 @@ def _lower_binary_unary(graph: Graph | GraphContext, node: Node) -> BinaryOp | U
         input0_shape = value_shape(graph, node.inputs[0], node)
         input1_shape = value_shape(graph, node.inputs[1], node)
         output_shape = value_shape(graph, node.outputs[0], node)
-        return BinaryOp(
+        op = BinaryOp(
             input0=node.inputs[0],
             input1=node.inputs[1],
             output=node.outputs[0],
             function=function,
             operator_kind=op_spec.kind,
-            input0_shape=input0_shape,
-            input1_shape=input1_shape,
-            shape=output_shape,
-            dtype=output_dtype,
-            input_dtype=input_dtype,
         )
+        if isinstance(graph, GraphContext):
+            inferred_shape = _infer_binary_output_shape(
+                function=function,
+                input0_shape=input0_shape,
+                input1_shape=input1_shape,
+            )
+            graph.set_shape(node.outputs[0], inferred_shape)
+        return op
     op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
     op_spec = binary_op_symbol(function, node.attrs, dtype=op_dtype)
     unary_symbol = unary_op_symbol(function, dtype=op_dtype)
@@ -226,32 +325,36 @@ def _lower_binary_unary(graph: Graph | GraphContext, node: Node) -> BinaryOp | U
         input0_shape = value_shape(graph, node.inputs[0], node)
         input1_shape = value_shape(graph, node.inputs[1], node)
         output_shape = value_shape(graph, node.outputs[0], node)
-        return BinaryOp(
+        op = BinaryOp(
             input0=node.inputs[0],
             input1=node.inputs[1],
             output=node.outputs[0],
             function=function,
             operator_kind=op_spec.kind,
-            input0_shape=input0_shape,
-            input1_shape=input1_shape,
-            shape=output_shape,
-            dtype=op_dtype,
-            input_dtype=op_dtype,
         )
+        if isinstance(graph, GraphContext):
+            inferred_shape = _infer_binary_output_shape(
+                function=function,
+                input0_shape=input0_shape,
+                input1_shape=input1_shape,
+            )
+            graph.set_shape(node.outputs[0], inferred_shape)
+        return op
     if len(node.inputs) != 1 or len(node.outputs) != 1:
         raise UnsupportedOpError(
             f"{node.op_type} must have 1 input and 1 output"
         )
     output_shape = value_shape(graph, node.outputs[0], node)
-    return UnaryOp(
+    op = UnaryOp(
         input0=node.inputs[0],
         output=node.outputs[0],
         function=function,
-        shape=output_shape,
-        dtype=op_dtype,
-        input_dtype=op_dtype,
         params=(),
     )
+    if isinstance(graph, GraphContext):
+        inferred_shape = value_shape(graph, node.inputs[0], node)
+        graph.set_shape(node.outputs[0], inferred_shape)
+    return op
 
 
 _DEFAULT_ELEMENTWISE_TYPES = (
@@ -283,9 +386,6 @@ def lower_isinf(graph: Graph, node: Node) -> UnaryOp:
         input0=node.inputs[0],
         output=node.outputs[0],
         function=ScalarFunction.ISINF,
-        shape=output_shape,
-        dtype=output_dtype,
-        input_dtype=input_dtype,
         params=(float(detect_negative), float(detect_positive)),
     )
 
@@ -305,8 +405,5 @@ def lower_isnan(graph: Graph, node: Node) -> UnaryOp:
         input0=node.inputs[0],
         output=node.outputs[0],
         function=ScalarFunction.ISNAN,
-        shape=output_shape,
-        dtype=output_dtype,
-        input_dtype=input_dtype,
         params=(),
     )

emx_onnx_cgen/lowering/gather.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
-from ..ir.ops import GatherOp
 from ..errors import UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
+from ..ir.ops import GatherOp
+from ..lowering.common import value_shape
+from ..validation import normalize_axis
 from .registry import register_lowering
 
 
@@ -11,9 +14,15 @@ def lower_gather(graph: Graph, node: Node) -> GatherOp:
     if len(node.inputs) != 2 or len(node.outputs) != 1:
         raise UnsupportedOpError("Gather must have 2 inputs and 1 output")
     data_name, indices_name = node.inputs
+    data_shape = value_shape(graph, data_name, node)
+    indices_shape = value_shape(graph, indices_name, node)
+    axis = normalize_axis(int(node.attrs.get("axis", 0)), data_shape, node)
+    output_shape = data_shape[:axis] + indices_shape + data_shape[axis + 1 :]
+    if isinstance(graph, GraphContext):
+        graph.set_shape(node.outputs[0], output_shape)
     return GatherOp(
         data=data_name,
         indices=indices_name,
         output=node.outputs[0],
-        axis=int(node.attrs.get("axis", 0)),
+        axis=axis,
     )

emx_onnx_cgen/lowering/gemm.py

@@ -1,139 +1,22 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-
-from shared.scalar_types import ScalarType
-
 from ..ir.ops import GemmOp
-from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
-from .common import node_dtype as _node_dtype
-from .common import value_shape as _value_shape
 from .registry import register_lowering
 
 
-@dataclass(frozen=True)
-class GemmSpec:
-    m: int
-    n: int
-    k: int
-    alpha: float | int
-    beta: float | int
-    trans_a: bool
-    trans_b: bool
-    c_shape: tuple[int, ...] | None
-
-
-def resolve_gemm_spec(graph: Graph, node: Node, dtype: ScalarType) -> GemmSpec:
-    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
-        raise UnsupportedOpError("Gemm must have 2 or 3 inputs and 1 output")
-    alpha, beta, trans_a, trans_b = _resolve_gemm_attrs(node, dtype)
-    input0_shape = _value_shape(graph, node.inputs[0], node)
-    input1_shape = _value_shape(graph, node.inputs[1], node)
-    if len(input0_shape) != 2 or len(input1_shape) != 2:
-        raise UnsupportedOpError(
-            "Gemm supports 2D inputs only, "
-            f"got {input0_shape} x {input1_shape}"
-        )
-    if trans_a:
-        m, k_left = input0_shape[1], input0_shape[0]
-    else:
-        m, k_left = input0_shape
-    if trans_b:
-        n, k_right = input1_shape[0], input1_shape[1]
-    else:
-        k_right, n = input1_shape
-    if k_left != k_right:
-        raise ShapeInferenceError(
-            f"Gemm inner dimensions must match, got {k_left} and {k_right}"
-        )
-    output_shape = _value_shape(graph, node.outputs[0], node)
-    if output_shape != (m, n):
-        raise ShapeInferenceError(
-            f"Gemm output shape must be {(m, n)}, got {output_shape}"
-        )
-    c_shape = None
-    if len(node.inputs) == 3:
-        bias_shape = _value_shape(graph, node.inputs[2], node)
-        c_shape = validate_gemm_bias_shape((m, n), bias_shape, node)
-    return GemmSpec(
-        m=m,
-        n=n,
-        k=k_left,
-        alpha=alpha,
-        beta=beta,
-        trans_a=trans_a,
-        trans_b=trans_b,
-        c_shape=c_shape,
-    )
-
-
-def _resolve_gemm_attrs(
-    node: Node, dtype: ScalarType
-) -> tuple[float | int, float | int, bool, bool]:
-    alpha = float(node.attrs.get("alpha", 1.0))
-    beta = float(node.attrs.get("beta", 1.0))
-    trans_a = int(node.attrs.get("transA", 0))
-    trans_b = int(node.attrs.get("transB", 0))
-    if trans_a not in {0, 1} or trans_b not in {0, 1}:
-        raise UnsupportedOpError(
-            "Gemm only supports transA/transB values of 0 or 1"
-        )
-    if dtype == ScalarType.BOOL:
-        raise UnsupportedOpError("Gemm supports numeric inputs only")
-    if not dtype.is_float:
-        alpha_int = int(alpha)
-        beta_int = int(beta)
-        if alpha != alpha_int or beta != beta_int:
-            raise UnsupportedOpError(
-                "Gemm alpha and beta must be integers for non-float inputs"
-            )
-        alpha = alpha_int
-        beta = beta_int
-    return alpha, beta, bool(trans_a), bool(trans_b)
-
-
-def validate_gemm_bias_shape(
-    output_shape: tuple[int, int], bias_shape: tuple[int, ...], node: Node
-) -> tuple[int, ...]:
-    if len(bias_shape) == 0:
-        return bias_shape
-    if len(bias_shape) == 1:
-        if bias_shape[0] not in {1, output_shape[1]}:
-            raise ShapeInferenceError(
-                "Gemm bias input must be broadcastable to output shape, "
-                f"got {bias_shape} vs {output_shape}"
-            )
-        return bias_shape
-    if len(bias_shape) == 2:
-        m, n = output_shape
-        if bias_shape[0] not in {1, m} or bias_shape[1] not in {1, n}:
-            raise ShapeInferenceError(
-                "Gemm bias input must be broadcastable to output shape, "
-                f"got {bias_shape} vs {output_shape}"
-            )
-        return bias_shape
-    raise ShapeInferenceError(
-        f"Gemm bias input must be rank 1 or 2, got {bias_shape}"
-    )
-
-
 @register_lowering("Gemm")
 def lower_gemm(graph: Graph, node: Node) -> GemmOp:
-    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
-    spec = resolve_gemm_spec(graph, node, op_dtype)
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError("Gemm must have 2 or 3 inputs and 1 output")
     return GemmOp(
         input_a=node.inputs[0],
         input_b=node.inputs[1],
         input_c=node.inputs[2] if len(node.inputs) == 3 else None,
         output=node.outputs[0],
-        m=spec.m,
-        n=spec.n,
-        k=spec.k,
-        trans_a=spec.trans_a,
-        trans_b=spec.trans_b,
-        alpha=spec.alpha,
-        beta=spec.beta,
-        c_shape=spec.c_shape,
-        dtype=op_dtype,
+        alpha=float(node.attrs.get("alpha", 1.0)),
+        beta=float(node.attrs.get("beta", 1.0)),
+        trans_a=int(node.attrs.get("transA", 0)),
+        trans_b=int(node.attrs.get("transB", 0)),
     )

emx_onnx_cgen/lowering/global_max_pool.py

@@ -45,15 +45,10 @@ def lower_global_max_pool(graph: Graph, node: Node) -> ReduceOp:
     return ReduceOp(
         input0=node.inputs[0],
         output=node.outputs[0],
-        input_shape=input_shape,
-        output_shape=output_shape,
         axes=axes,
         axes_input=None,
-        axes_input_shape=None,
-        axes_input_dtype=None,
         keepdims=True,
         noop_with_empty_axes=False,
         reduce_kind="max",
         reduce_count=None,
-        dtype=op_dtype,
     )