emx-onnx-cgen 0.3.8__py3-none-any.whl → 0.4.2.dev0__py3-none-any.whl

emx_onnx_cgen/ir/ops/nn.py

@@ -6,7 +6,7 @@ from enum import Enum
 from shared.scalar_functions import ScalarFunction
 from shared.scalar_types import ScalarType
 
-from ...errors import ShapeInferenceError
+from ...errors import ShapeInferenceError, UnsupportedOpError
 from ..op_base import ConvLikeOpBase, GemmLikeOpBase, MatMulLikeOpBase, RenderableOpBase
 from ..op_context import OpContext
 
@@ -29,23 +29,117 @@ def _shape_product(shape: tuple[int, ...]) -> int:
     return product
 
 
+def _broadcast_batch_shapes(
+    left: tuple[int, ...], right: tuple[int, ...]
+) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+    max_rank = max(len(left), len(right))
+    left_padded = (1,) * (max_rank - len(left)) + left
+    right_padded = (1,) * (max_rank - len(right)) + right
+    broadcast_shape: list[int] = []
+    for left_dim, right_dim in zip(left_padded, right_padded):
+        if not (left_dim == right_dim or left_dim == 1 or right_dim == 1):
+            raise ShapeInferenceError(
+                "MatMul batch dimensions must be broadcastable, "
+                f"got {left} x {right}"
+            )
+        broadcast_shape.append(max(left_dim, right_dim))
+    return tuple(broadcast_shape), left_padded, right_padded
+
+
+def _resolve_matmul_spec(
+    ctx: OpContext, input0: str, input1: str
+) -> dict[str, object]:
+    input0_shape = ctx.shape(input0)
+    input1_shape = ctx.shape(input1)
+    if len(input0_shape) < 1 or len(input1_shape) < 1:
+        raise UnsupportedOpError(
+            "MatMul inputs must be at least 1D, "
+            f"got {input0_shape} x {input1_shape}"
+        )
+    left_vector = len(input0_shape) == 1
+    right_vector = len(input1_shape) == 1
+    input0_effective = (1, input0_shape[0]) if left_vector else input0_shape
+    input1_effective = (input1_shape[0], 1) if right_vector else input1_shape
+    m, k_left = input0_effective[-2], input0_effective[-1]
+    k_right, n = input1_effective[-2], input1_effective[-1]
+    if k_left != k_right:
+        raise ShapeInferenceError(
+            f"MatMul inner dimensions must match, got {k_left} and {k_right}"
+        )
+    batch_shape, input0_batch_shape, input1_batch_shape = (
+        _broadcast_batch_shapes(
+            input0_effective[:-2],
+            input1_effective[:-2],
+        )
+    )
+    if left_vector and right_vector:
+        output_shape = batch_shape
+    elif left_vector:
+        output_shape = batch_shape + (n,)
+    elif right_vector:
+        output_shape = batch_shape + (m,)
+    else:
+        output_shape = batch_shape + (m, n)
+    return {
+        "input0_shape": input0_shape,
+        "input1_shape": input1_shape,
+        "output_shape": output_shape,
+        "batch_shape": batch_shape,
+        "input0_batch_shape": input0_batch_shape,
+        "input1_batch_shape": input1_batch_shape,
+        "m": m,
+        "n": n,
+        "k": k_left,
+        "left_vector": left_vector,
+        "right_vector": right_vector,
+    }
+
+
 @dataclass(frozen=True)
 class MatMulOp(MatMulLikeOpBase):
     input0: str
     input1: str
     output: str
-    input0_shape: tuple[int, ...]
-    input1_shape: tuple[int, ...]
-    output_shape: tuple[int, ...]
-    batch_shape: tuple[int, ...]
-    input0_batch_shape: tuple[int, ...]
-    input1_batch_shape: tuple[int, ...]
-    m: int
-    n: int
-    k: int
-    left_vector: bool
-    right_vector: bool
-    dtype: ScalarType
+
+    def infer_types(self, ctx: OpContext) -> None:
+        input0_dtype = ctx.dtype(self.input0)
+        input1_dtype = ctx.dtype(self.input1)
+        if input0_dtype != input1_dtype:
+            raise UnsupportedOpError(
+                "MatMul expects matching input dtypes, "
+                f"got {input0_dtype.onnx_name} and {input1_dtype.onnx_name}"
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, input0_dtype)
+            output_dtype = input0_dtype
+        if output_dtype != input0_dtype:
+            raise UnsupportedOpError(
+                "MatMul expects output dtype to match inputs, "
+                f"got {output_dtype.onnx_name} and {input0_dtype.onnx_name}"
+            )
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        spec = _resolve_matmul_spec(ctx, self.input0, self.input1)
+        output_shape = spec["output_shape"]
+        try:
+            expected = ctx.shape(self.output)
+        except ShapeInferenceError:
+            expected = None
+        if expected is not None and expected != output_shape:
+            raise ShapeInferenceError(
+                f"MatMul output shape must be {output_shape}, got {expected}"
+            )
+        ctx.set_shape(self.output, output_shape)
+        ctx.set_derived(self, "batch_shape", spec["batch_shape"])
+        ctx.set_derived(self, "input0_batch_shape", spec["input0_batch_shape"])
+        ctx.set_derived(self, "input1_batch_shape", spec["input1_batch_shape"])
+        ctx.set_derived(self, "m", spec["m"])
+        ctx.set_derived(self, "n", spec["n"])
+        ctx.set_derived(self, "k", spec["k"])
+        ctx.set_derived(self, "left_vector", spec["left_vector"])
+        ctx.set_derived(self, "right_vector", spec["right_vector"])
 
 @dataclass(frozen=True)
 class QLinearMatMulOp(MatMulLikeOpBase):
@@ -98,15 +192,139 @@ class GemmOp(GemmLikeOpBase):
     input_b: str
     input_c: str | None
     output: str
-    m: int
-    n: int
-    k: int
-    trans_a: bool
-    trans_b: bool
+    trans_a: int
+    trans_b: int
     alpha: float | int
     beta: float | int
-    c_shape: tuple[int, ...] | None
-    dtype: ScalarType
+
+    @staticmethod
+    def _normalize_attrs(
+        dtype: ScalarType,
+        *,
+        alpha: float | int,
+        beta: float | int,
+        trans_a: int,
+        trans_b: int,
+    ) -> tuple[float | int, float | int, bool, bool]:
+        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)
+
+    @staticmethod
+    def _validate_bias_shape(
+        output_shape: tuple[int, int], bias_shape: tuple[int, ...]
+    ) -> 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}"
+        )
+
+    def infer_types(self, ctx: OpContext) -> None:
+        input_a_dtype = ctx.dtype(self.input_a)
+        input_b_dtype = ctx.dtype(self.input_b)
+        if input_a_dtype != input_b_dtype:
+            raise UnsupportedOpError(
+                "Gemm expects matching input dtypes, "
+                f"got {input_a_dtype.onnx_name} and {input_b_dtype.onnx_name}"
+            )
+        if self.input_c is not None:
+            input_c_dtype = ctx.dtype(self.input_c)
+            if input_c_dtype != input_a_dtype:
+                raise UnsupportedOpError(
+                    "Gemm expects bias dtype to match inputs, "
+                    f"got {input_c_dtype.onnx_name} and {input_a_dtype.onnx_name}"
+                )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, input_a_dtype)
+            output_dtype = input_a_dtype
+        if output_dtype != input_a_dtype:
+            raise UnsupportedOpError(
+                "Gemm expects output dtype to match inputs, "
+                f"got {output_dtype.onnx_name} and {input_a_dtype.onnx_name}"
+            )
+        alpha, beta, trans_a, trans_b = self._normalize_attrs(
+            output_dtype,
+            alpha=self.alpha,
+            beta=self.beta,
+            trans_a=self.trans_a,
+            trans_b=self.trans_b,
+        )
+        ctx.set_derived(self, "alpha", alpha)
+        ctx.set_derived(self, "beta", beta)
+        ctx.set_derived(self, "trans_a", trans_a)
+        ctx.set_derived(self, "trans_b", trans_b)
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        trans_a = ctx.require_derived(self, "trans_a")
+        trans_b = ctx.require_derived(self, "trans_b")
+        input_a_shape = ctx.shape(self.input_a)
+        input_b_shape = ctx.shape(self.input_b)
+        if len(input_a_shape) != 2 or len(input_b_shape) != 2:
+            raise UnsupportedOpError(
+                "Gemm supports 2D inputs only, "
+                f"got {input_a_shape} x {input_b_shape}"
+            )
+        if trans_a:
+            m, k_left = input_a_shape[1], input_a_shape[0]
+        else:
+            m, k_left = input_a_shape
+        if trans_b:
+            n, k_right = input_b_shape[0], input_b_shape[1]
+        else:
+            k_right, n = input_b_shape
+        if k_left != k_right:
+            raise ShapeInferenceError(
+                f"Gemm inner dimensions must match, got {k_left} and {k_right}"
+            )
+        output_shape = (m, n)
+        try:
+            expected = ctx.shape(self.output)
+        except ShapeInferenceError:
+            expected = None
+        if expected is not None and expected != output_shape:
+            raise ShapeInferenceError(
+                f"Gemm output shape must be {output_shape}, got {expected}"
+            )
+        ctx.set_shape(self.output, output_shape)
+        c_shape = None
+        if self.input_c is not None:
+            bias_shape = ctx.shape(self.input_c)
+            c_shape = self._validate_bias_shape(output_shape, bias_shape)
+        ctx.set_derived(self, "m", m)
+        ctx.set_derived(self, "n", n)
+        ctx.set_derived(self, "k", k_left)
+        ctx.set_derived(self, "c_shape", c_shape)
 
 @dataclass(frozen=True)
 class AttentionOp(RenderableOpBase):
@@ -205,6 +423,31 @@ class ConvOp(ConvLikeOpBase):
             raise ValueError("Conv output width is undefined for spatial_rank < 2")
         return self.out_spatial[1]
 
+@dataclass(frozen=True)
+class ConvIntegerOp(ConvLikeOpBase):
+    input0: str
+    weights: str
+    x_zero_point: str | None
+    w_zero_point: str | None
+    output: str
+    batch: int
+    in_channels: int
+    out_channels: int
+    spatial_rank: int
+    in_spatial: tuple[int, ...]
+    out_spatial: tuple[int, ...]
+    kernel_shape: tuple[int, ...]
+    strides: tuple[int, ...]
+    pads: tuple[int, ...]
+    dilations: tuple[int, ...]
+    group: int
+    input_dtype: ScalarType
+    weight_dtype: ScalarType
+    dtype: ScalarType
+    x_zero_point_shape: tuple[int, ...] | None
+    w_zero_point_shape: tuple[int, ...] | None
+    w_zero_point_per_channel: bool
+
 @dataclass(frozen=True)
 class ConvTransposeOp(ConvLikeOpBase):
     input0: str
@@ -237,6 +480,8 @@ class AveragePoolOp(RenderableOpBase):
     out_w: int
     kernel_h: int
     kernel_w: int
+    dilation_h: int
+    dilation_w: int
     stride_h: int
     stride_w: int
     pad_top: int
@@ -245,6 +490,14 @@ class AveragePoolOp(RenderableOpBase):
     pad_right: int
     count_include_pad: bool
     dtype: ScalarType
+    spatial_rank: int = 2
+    in_d: int = 1
+    out_d: int = 1
+    kernel_d: int = 1
+    dilation_d: int = 1
+    stride_d: int = 1
+    pad_front: int = 0
+    pad_back: int = 0
 
 @dataclass(frozen=True)
 class LpPoolOp(RenderableOpBase):
@@ -258,6 +511,8 @@ class LpPoolOp(RenderableOpBase):
     out_w: int
     kernel_h: int
     kernel_w: int
+    dilation_h: int
+    dilation_w: int
     stride_h: int
     stride_w: int
     pad_top: int
@@ -271,16 +526,29 @@ class LpPoolOp(RenderableOpBase):
 class SoftmaxOp(RenderableOpBase):
     input0: str
     output: str
-    outer: int
-    axis_size: int
-    inner: int
-    axis: int
-    shape: tuple[int, ...]
-    dtype: ScalarType
+    axis: int | None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        input_dtype = ctx.dtype(self.input0)
+        if not input_dtype.is_float:
+            raise UnsupportedOpError(
+                "Softmax supports float16, float, and double inputs only"
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, input_dtype)
+            return None
+        if output_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "Softmax expects output dtype to match input dtype"
+            )
 
     def infer_shapes(self, ctx: OpContext) -> None:
         input_shape = ctx.shape(self.input0)
         axis = self.axis
+        if axis is None:
+            axis = -1
         if axis < 0:
             axis += len(input_shape)
         if axis < 0 or axis >= len(input_shape):
@@ -295,6 +563,7 @@ class SoftmaxOp(RenderableOpBase):
             else 1
         )
         ctx.set_shape(self.output, input_shape)
+        ctx.set_derived(self, "axis", axis)
         ctx.set_derived(self, "outer", outer)
         ctx.set_derived(self, "axis_size", axis_size)
         ctx.set_derived(self, "inner", inner)
@@ -303,16 +572,29 @@ class SoftmaxOp(RenderableOpBase):
 class LogSoftmaxOp(RenderableOpBase):
     input0: str
     output: str
-    outer: int
-    axis_size: int
-    inner: int
-    axis: int
-    shape: tuple[int, ...]
-    dtype: ScalarType
+    axis: int | None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        input_dtype = ctx.dtype(self.input0)
+        if not input_dtype.is_float:
+            raise UnsupportedOpError(
+                "LogSoftmax supports float16, float, and double inputs only"
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, input_dtype)
+            return None
+        if output_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "LogSoftmax expects output dtype to match input dtype"
+            )
 
     def infer_shapes(self, ctx: OpContext) -> None:
         input_shape = ctx.shape(self.input0)
         axis = self.axis
+        if axis is None:
+            axis = -1
         if axis < 0:
             axis += len(input_shape)
         if axis < 0 or axis >= len(input_shape):
@@ -327,6 +609,7 @@ class LogSoftmaxOp(RenderableOpBase):
             else 1
         )
         ctx.set_shape(self.output, input_shape)
+        ctx.set_derived(self, "axis", axis)
         ctx.set_derived(self, "outer", outer)
         ctx.set_derived(self, "axis_size", axis_size)
         ctx.set_derived(self, "inner", inner)
@@ -335,16 +618,30 @@ class LogSoftmaxOp(RenderableOpBase):
 class HardmaxOp(RenderableOpBase):
     input0: str
     output: str
-    outer: int
-    axis_size: int
-    inner: int
-    axis: int
-    shape: tuple[int, ...]
-    dtype: ScalarType
+    axis: int | None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        input_dtype = ctx.dtype(self.input0)
+        if input_dtype not in {ScalarType.F16, ScalarType.F32, ScalarType.F64}:
+            raise UnsupportedOpError(
+                "Hardmax supports float16, float, and double inputs only"
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, input_dtype)
+            return None
+        if output_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "Hardmax expects output dtype to match input dtype"
+            )
 
     def infer_shapes(self, ctx: OpContext) -> None:
         input_shape = ctx.shape(self.input0)
         axis = self.axis
+        if axis is None:
+            opset_version = ctx.opset_version()
+            axis = 1 if opset_version is not None and opset_version < 13 else -1
         if axis < 0:
             axis += len(input_shape)
         if axis < 0 or axis >= len(input_shape):
@@ -359,6 +656,7 @@ class HardmaxOp(RenderableOpBase):
             else 1
         )
         ctx.set_shape(self.output, input_shape)
+        ctx.set_derived(self, "axis", axis)
         ctx.set_derived(self, "outer", outer)
         ctx.set_derived(self, "axis_size", axis_size)
         ctx.set_derived(self, "inner", inner)
@@ -512,6 +810,31 @@ class LrnOp(RenderableOpBase):
     bias: float
     dtype: ScalarType
 
+@dataclass(frozen=True)
+class GruOp(RenderableOpBase):
+    input_x: str
+    input_w: str
+    input_r: str
+    input_b: str | None
+    input_sequence_lens: str | None
+    input_initial_h: str | None
+    output_y: str | None
+    output_y_h: str | None
+    seq_length: int
+    batch_size: int
+    input_size: int
+    hidden_size: int
+    num_directions: int
+    direction: str
+    layout: int
+    linear_before_reset: int
+    clip: float | None
+    activation_kinds: tuple[int, ...]
+    activation_alphas: tuple[float, ...]
+    activation_betas: tuple[float, ...]
+    dtype: ScalarType
+    sequence_lens_dtype: ScalarType | None
+
 @dataclass(frozen=True)
 class LstmOp(RenderableOpBase):
     input_x: str

emx_onnx_cgen/ir/ops/reduce.py

@@ -2,8 +2,6 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from shared.scalar_types import ScalarType
-
 from ..op_base import ReduceOpBase
 from ..op_context import OpContext
 
@@ -12,17 +10,12 @@ from ..op_context import OpContext
 class ReduceOp(ReduceOpBase):
     input0: str
     output: str
-    input_shape: tuple[int, ...]
-    output_shape: tuple[int, ...]
     axes: tuple[int, ...]
     axes_input: str | None
-    axes_input_shape: tuple[int, ...] | None
-    axes_input_dtype: ScalarType | None
     keepdims: bool
     noop_with_empty_axes: bool
     reduce_kind: str
     reduce_count: int | None
-    dtype: ScalarType
 
     def infer_types(self, ctx: OpContext) -> None:
         ctx.dtype(self.output)
@@ -45,14 +38,10 @@ class ReduceOp(ReduceOpBase):
 class ArgReduceOp(ReduceOpBase):
     input0: str
     output: str
-    input_shape: tuple[int, ...]
-    output_shape: tuple[int, ...]
     axis: int
     keepdims: bool
     select_last_index: bool
     reduce_kind: str
-    input_dtype: ScalarType
-    output_dtype: ScalarType
 
     def infer_types(self, ctx: OpContext) -> None:
         ctx.dtype(self.input0)
@@ -71,17 +60,13 @@ class ArgReduceOp(ReduceOpBase):
 @dataclass(frozen=True)
 class TopKOp(ReduceOpBase):
     input0: str
+    k_input: str
     output_values: str
     output_indices: str
-    input_shape: tuple[int, ...]
-    output_shape: tuple[int, ...]
     axis: int
     k: int
     largest: bool
     sorted: bool
-    input_dtype: ScalarType
-    output_values_dtype: ScalarType
-    output_indices_dtype: ScalarType
 
     def infer_types(self, ctx: OpContext) -> None:
         ctx.dtype(self.input0)

emx_onnx_cgen/lowering/__init__.py

@@ -10,13 +10,16 @@ _LOWERING_MODULES = [
     "attention",
     "average_pool",
     "batch_normalization",
+    "bernoulli",
     "cast",
     "concat",
     "constant_of_shape",
     "conv",
+    "conv_integer",
     "conv_transpose",
     "cumsum",
     "depth_space",
+    "dequantize_linear",
     "dropout",
     "einsum",
     "elementwise",
@@ -29,8 +32,10 @@ _LOWERING_MODULES = [
     "gemm",
     "global_max_pool",
     "grid_sample",
+    "gru",
     "group_normalization",
     "hardmax",
+    "hamming_window",
     "identity",
     "instance_normalization",
     "layer_normalization",
@@ -45,9 +50,11 @@ _LOWERING_MODULES = [
     "negative_log_likelihood_loss",
     "non_max_suppression",
     "nonzero",
+    "optional_has_element",
     "one_hot",
     "pad",
     "qlinear_matmul",
+    "qlinear_mul",
     "quantize_linear",
     "range",
     "reduce",
@@ -64,11 +71,13 @@ _LOWERING_MODULES = [
     "split",
     "squeeze",
     "tensor_scatter",
+    "tfidf_vectorizer",
     "tile",
     "topk",
     "transpose",
     "trilu",
     "unsqueeze",
+    "upsample",
     "variadic",
     "where",
 ]

emx_onnx_cgen/lowering/arg_reduce.py

@@ -84,14 +84,10 @@ def lower_arg_reduce(graph: Graph, node: Node) -> ArgReduceOp:
     return ArgReduceOp(
         input0=input_name,
         output=output_name,
-        input_shape=input_shape,
-        output_shape=output_shape,
         axis=axis,
         keepdims=keepdims,
         select_last_index=select_last_index,
         reduce_kind=ARG_REDUCE_KIND_BY_OP[node.op_type],
-        input_dtype=input_dtype,
-        output_dtype=output_dtype,
     )