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

emx_onnx_cgen/ir/ops/elementwise.py

@@ -6,7 +6,13 @@ from shared.scalar_functions import ScalarFunction
 from shared.scalar_types import ScalarType
 
 from ...ops import COMPARE_FUNCTIONS, OperatorKind, binary_op_symbol
-from ..op_base import ElementwiseOpBase, VariadicLikeOpBase
+from ...errors import ShapeInferenceError, UnsupportedOpError
+from ..op_base import (
+    BroadcastingOpBase,
+    ElementwiseOpBase,
+    RenderableOpBase,
+    VariadicLikeOpBase,
+)
 from ..op_context import OpContext
 
 
@@ -17,11 +23,6 @@ class BinaryOp(ElementwiseOpBase):
     output: str
     function: ScalarFunction
     operator_kind: OperatorKind
-    input0_shape: tuple[int, ...]
-    input1_shape: tuple[int, ...]
-    shape: tuple[int, ...]
-    dtype: ScalarType
-    input_dtype: ScalarType
 
     def _elementwise_inputs(self) -> tuple[str, ...]:
         return (self.input0, self.input1)
@@ -32,6 +33,107 @@ class BinaryOp(ElementwiseOpBase):
     def _elementwise_compare(self) -> bool:
         return self.function in COMPARE_FUNCTIONS
 
+    def infer_shapes(self, ctx: OpContext) -> None:
+        if self.function != ScalarFunction.PRELU:
+            return super().infer_shapes(ctx)
+        input_shape = ctx.shape(self.input0)
+        slope_shape = ctx.shape(self.input1)
+        output_name = self.output
+        if BroadcastingOpBase.unidirectional_broadcastable(
+            slope_shape, input_shape
+        ):
+            ctx.set_shape(output_name, input_shape)
+            return None
+        channel_axis = BroadcastingOpBase.prelu_channel_axis(
+            input_shape, slope_shape
+        )
+        if channel_axis is not None:
+            ctx.set_shape(output_name, input_shape)
+            ctx.set_derived(self, "prelu_slope_axis", channel_axis)
+            return None
+        raise ShapeInferenceError(
+            "Broadcasting mismatch for shapes: "
+            + ", ".join(str(shape) for shape in (input_shape, slope_shape))
+        )
+
+
+_POW_BASE_DTYPES = {
+    ScalarType.F16,
+    ScalarType.F32,
+    ScalarType.F64,
+    ScalarType.I32,
+    ScalarType.I64,
+}
+_POW_EXPONENT_DTYPES = {
+    ScalarType.F16,
+    ScalarType.F32,
+    ScalarType.F64,
+    ScalarType.I8,
+    ScalarType.I16,
+    ScalarType.I32,
+    ScalarType.I64,
+    ScalarType.U8,
+    ScalarType.U16,
+    ScalarType.U32,
+    ScalarType.U64,
+}
+
+
+@dataclass(frozen=True)
+class PowOp(BinaryOp):
+    def validate(self, ctx: OpContext) -> None:
+        base_dtype = ctx.dtype(self.input0)
+        exponent_dtype = ctx.dtype(self.input1)
+        if base_dtype not in _POW_BASE_DTYPES:
+            raise UnsupportedOpError(
+                "Pow base dtype must be one of "
+                f"{', '.join(dtype.onnx_name for dtype in sorted(_POW_BASE_DTYPES, key=str))}, "
+                f"got {base_dtype.onnx_name}"
+            )
+        if exponent_dtype not in _POW_EXPONENT_DTYPES:
+            raise UnsupportedOpError(
+                "Pow exponent dtype must be one of "
+                f"{', '.join(dtype.onnx_name for dtype in sorted(_POW_EXPONENT_DTYPES, key=str))}, "
+                f"got {exponent_dtype.onnx_name}"
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            return None
+        if output_dtype != base_dtype:
+            raise UnsupportedOpError(
+                "Pow expects output dtype "
+                f"{base_dtype.onnx_name}, got {output_dtype.onnx_name}"
+            )
+        return None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        base_dtype = ctx.dtype(self.input0)
+        exponent_dtype = ctx.dtype(self.input1)
+        if base_dtype not in _POW_BASE_DTYPES:
+            raise UnsupportedOpError(
+                "Pow base dtype must be one of "
+                f"{', '.join(dtype.onnx_name for dtype in sorted(_POW_BASE_DTYPES, key=str))}, "
+                f"got {base_dtype.onnx_name}"
+            )
+        if exponent_dtype not in _POW_EXPONENT_DTYPES:
+            raise UnsupportedOpError(
+                "Pow exponent dtype must be one of "
+                f"{', '.join(dtype.onnx_name for dtype in sorted(_POW_EXPONENT_DTYPES, key=str))}, "
+                f"got {exponent_dtype.onnx_name}"
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, base_dtype)
+            return None
+        if output_dtype != base_dtype:
+            raise UnsupportedOpError(
+                "Pow expects output dtype "
+                f"{base_dtype.onnx_name}, got {output_dtype.onnx_name}"
+            )
+        return None
+
 
 @dataclass(frozen=True)
 class VariadicOp(VariadicLikeOpBase):
@@ -80,11 +182,6 @@ class WhereOp(ElementwiseOpBase):
     input_x: str
     input_y: str
     output: str
-    condition_shape: tuple[int, ...]
-    x_shape: tuple[int, ...]
-    y_shape: tuple[int, ...]
-    output_shape: tuple[int, ...]
-    dtype: ScalarType
 
     def _elementwise_inputs(self) -> tuple[str, ...]:
         return (self.condition, self.input_x, self.input_y)
@@ -101,9 +198,6 @@ class UnaryOp(ElementwiseOpBase):
     input0: str
     output: str
     function: ScalarFunction
-    shape: tuple[int, ...]
-    dtype: ScalarType
-    input_dtype: ScalarType
     params: tuple[float, ...] = ()
 
     def _elementwise_inputs(self) -> tuple[str, ...]:
@@ -126,11 +220,8 @@ class ClipOp(ElementwiseOpBase):
     input_min: str | None
     input_max: str | None
     output: str
-    input_shape: tuple[int, ...]
-    min_shape: tuple[int, ...] | None
-    max_shape: tuple[int, ...] | None
-    output_shape: tuple[int, ...]
-    dtype: ScalarType
+    min_value: float | None = None
+    max_value: float | None = None
 
     def _elementwise_inputs(self) -> tuple[str, ...]:
         inputs = [self.input0]
@@ -152,9 +243,6 @@ class ClipOp(ElementwiseOpBase):
 class IdentityOp(ElementwiseOpBase):
     input0: str
     output: str
-    shape: tuple[int, ...]
-    dtype: ScalarType
-    input_dtype: ScalarType
 
     def _elementwise_inputs(self) -> tuple[str, ...]:
         return (self.input0,)
@@ -165,3 +253,31 @@ class IdentityOp(ElementwiseOpBase):
     def validate(self, ctx: OpContext) -> None:
         super().validate(ctx)
         return None
+
+
+@dataclass(frozen=True)
+class QLinearMulOp(RenderableOpBase):
+    input0: str
+    input0_scale: str
+    input0_zero_point: str
+    input1: str
+    input1_scale: str
+    input1_zero_point: str
+    output_scale: str
+    output_zero_point: str
+    output: str
+    input0_shape: tuple[int, ...]
+    input1_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    input0_dtype: ScalarType
+    input1_dtype: ScalarType
+    dtype: ScalarType
+    input0_scale_dtype: ScalarType
+    input1_scale_dtype: ScalarType
+    output_scale_dtype: ScalarType
+    input0_scale_shape: tuple[int, ...]
+    input1_scale_shape: tuple[int, ...]
+    output_scale_shape: tuple[int, ...]
+    input0_zero_shape: tuple[int, ...]
+    input1_zero_shape: tuple[int, ...]
+    output_zero_shape: tuple[int, ...]

emx_onnx_cgen/ir/ops/misc.py

@@ -53,6 +53,20 @@ class QuantizeLinearOp(RenderableOpBase):
     input_dtype: ScalarType
     scale_dtype: ScalarType
 
+
+@dataclass(frozen=True)
+class DequantizeLinearOp(RenderableOpBase):
+    input0: str
+    scale: str
+    zero_point: str | None
+    output: str
+    input_shape: tuple[int, ...]
+    axis: int | None
+    block_size: int | None
+    dtype: ScalarType
+    input_dtype: ScalarType
+    scale_dtype: ScalarType
+
 @dataclass(frozen=True)
 class ConcatOp(RenderableOpBase):
     inputs: tuple[str, ...]
@@ -184,6 +198,16 @@ class EyeLikeOp(RenderableOpBase):
     dtype: ScalarType
     input_dtype: ScalarType
 
+@dataclass(frozen=True)
+class BernoulliOp(RenderableOpBase):
+    input0: str
+    output: str
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    input_dtype: ScalarType
+    dtype: ScalarType
+    seed: int | None
+
 @dataclass(frozen=True)
 class TriluOp(RenderableOpBase):
     input0: str
@@ -355,6 +379,51 @@ class SizeOp(RenderableOpBase):
     dtype: ScalarType
     input_dtype: ScalarType
 
+
+@dataclass(frozen=True)
+class OptionalHasElementOp(RenderableOpBase):
+    input0: str
+    output: str
+
+    def validate(self, ctx: OpContext) -> None:
+        value = ctx.graph.find_value(self.input0)
+        if not value.type.is_optional:
+            raise UnsupportedOpError(
+                f"{self.kind} expects optional input, got non-optional tensor."
+            )
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            return None
+        if output_dtype != ScalarType.BOOL:
+            raise UnsupportedOpError(
+                f"{self.kind} expects bool output, got {output_dtype.onnx_name}"
+            )
+        return None
+
+    def infer_types(self, ctx: OpContext) -> None:
+        ctx.dtype(self.input0)
+        try:
+            output_dtype = ctx.dtype(self.output)
+        except ShapeInferenceError:
+            ctx.set_dtype(self.output, ScalarType.BOOL)
+            return None
+        if output_dtype != ScalarType.BOOL:
+            raise UnsupportedOpError(
+                f"{self.kind} expects bool output, got {output_dtype.onnx_name}"
+            )
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        try:
+            output_shape = ctx.shape(self.output)
+        except ShapeInferenceError:
+            ctx.set_shape(self.output, ())
+            return None
+        if output_shape not in {(), (1,)}:
+            raise UnsupportedOpError(
+                f"{self.kind} expects scalar output, got shape {output_shape}"
+            )
+
 @dataclass(frozen=True)
 class NonZeroOp(RenderableOpBase):
     input0: str
@@ -477,6 +546,15 @@ class RangeOp(RenderableOpBase):
     dtype: ScalarType
     input_dtype: ScalarType
 
+@dataclass(frozen=True)
+class HammingWindowOp(RenderableOpBase):
+    size: str
+    output: str
+    output_shape: tuple[int, ...]
+    periodic: bool
+    dtype: ScalarType
+    input_dtype: ScalarType
+
 @dataclass(frozen=True)
 class OneHotOp(RenderableOpBase):
     indices: str
@@ -492,6 +570,23 @@ class OneHotOp(RenderableOpBase):
     indices_dtype: ScalarType
     depth_dtype: ScalarType
 
+@dataclass(frozen=True)
+class TfIdfVectorizerOp(RenderableOpBase):
+    input0: str
+    output: str
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    input_dtype: ScalarType
+    output_dtype: ScalarType
+    min_gram_length: int
+    max_gram_length: int
+    max_skip_count: int
+    mode: str
+    ngram_counts: tuple[int, ...]
+    ngram_indexes: tuple[int, ...]
+    pool_int64s: tuple[int, ...]
+    weights: tuple[float, ...] | None
+
 @dataclass(frozen=True)
 class SplitOp(RenderableOpBase):
     input0: str