emx-onnx-cgen 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

emx_onnx_cgen/lowering/softmax.py

@@ -4,6 +4,7 @@ from ..codegen.c_emitter import SoftmaxOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype
+from .common import onnx_opset_version as _onnx_opset_version
 from .common import shape_product as _shape_product
 from .common import value_shape as _value_shape
 from .registry import register_lowering
@@ -23,8 +24,11 @@ def lower_softmax(graph: Graph, node: Node) -> SoftmaxOp:
     input_shape = _value_shape(graph, node.inputs[0], node)
     output_shape = _value_shape(graph, node.outputs[0], node)
     ensure_output_shape_matches_input(node, input_shape, output_shape)
+    opset_version = _onnx_opset_version(graph)
+    default_axis = 1 if opset_version is not None and opset_version < 13 else -1
+    axis_attr = node.attrs.get("axis", default_axis)
     axis = _normalize_axis(
-        int(node.attrs.get("axis", -1)),
+        int(axis_attr),
         input_shape,
         node,
     )

emx_onnx_cgen/lowering/squeeze.py

@@ -95,11 +95,11 @@ def _validate_output_shape_for_unknown_axes(
     for dim in input_shape:
         if output_index < len(output_shape) and dim == output_shape[output_index]:
             output_index += 1
-            continue
-        if dim != 1:
-            raise ShapeInferenceError(
-                "Squeeze output shape must remove only dimensions of size 1"
-            )
+        else:
+            if dim != 1:
+                raise ShapeInferenceError(
+                    "Squeeze output shape must remove only dimensions of size 1"
+                )
     if output_index != len(output_shape):
         raise ShapeInferenceError(
             "Squeeze output shape must preserve input order while removing size-1 axes"

emx_onnx_cgen/lowering/topk.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import TopKOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import shape_product, value_dtype, value_shape
+from ..validation import normalize_axis
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _read_k(graph: Graph, name: str, node: Node) -> int:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        raise UnsupportedOpError(
+            f"{node.op_type} k input must be a constant initializer"
+        )
+    if initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            f"{node.op_type} k input must be int64 or int32"
+        )
+    data = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if data.size != 1:
+        raise ShapeInferenceError(
+            f"{node.op_type} k input must contain a single value"
+        )
+    k = int(data[0])
+    if k <= 0:
+        raise ShapeInferenceError(
+            f"{node.op_type} k must be a positive value, got {k}"
+        )
+    return k
+
+
+def _topk_dtype_supported(dtype: ScalarType) -> bool:
+    return not dtype.is_bool
+
+
+def lower_topk(graph: Graph, node: Node) -> TopKOp:
+    if node.op_type != "TopK":
+        raise UnsupportedOpError(f"Unsupported op {node.op_type}")
+    if len(node.inputs) != 2 or len(node.outputs) != 2:
+        raise UnsupportedOpError(
+            f"{node.op_type} must have 2 inputs and 2 outputs"
+        )
+    input_name = node.inputs[0]
+    k_name = node.inputs[1]
+    output_values = node.outputs[0]
+    output_indices = node.outputs[1]
+    input_shape = value_shape(graph, input_name, node)
+    shape_product(input_shape)
+    axis = int(node.attrs.get("axis", -1))
+    axis = normalize_axis(axis, input_shape, node)
+    k = _read_k(graph, k_name, node)
+    axis_dim = input_shape[axis]
+    if k > axis_dim:
+        raise ShapeInferenceError(
+            f"{node.op_type} k {k} exceeds axis dimension {axis_dim}"
+        )
+    output_shape_expected = list(input_shape)
+    output_shape_expected[axis] = k
+    output_shape = tuple(output_shape_expected)
+    values_shape = value_shape(graph, output_values, node)
+    if values_shape != output_shape:
+        raise ShapeInferenceError(
+            f"{node.op_type} values output shape must be {output_shape}, got {values_shape}"
+        )
+    indices_shape = value_shape(graph, output_indices, node)
+    if indices_shape != output_shape:
+        raise ShapeInferenceError(
+            f"{node.op_type} indices output shape must be {output_shape}, got {indices_shape}"
+        )
+    input_dtype = value_dtype(graph, input_name, node)
+    if not _topk_dtype_supported(input_dtype):
+        raise UnsupportedOpError(
+            f"{node.op_type} does not support dtype {input_dtype.onnx_name}"
+        )
+    values_dtype = value_dtype(graph, output_values, node)
+    if values_dtype != input_dtype:
+        raise UnsupportedOpError(
+            f"{node.op_type} values output dtype must be {input_dtype.onnx_name}"
+        )
+    indices_dtype = value_dtype(graph, output_indices, node)
+    if indices_dtype != ScalarType.I64:
+        raise UnsupportedOpError(
+            f"{node.op_type} indices output dtype must be int64"
+        )
+    largest = bool(int(node.attrs.get("largest", 1)))
+    sorted_output = bool(int(node.attrs.get("sorted", 1)))
+    return TopKOp(
+        input0=input_name,
+        output_values=output_values,
+        output_indices=output_indices,
+        input_shape=input_shape,
+        output_shape=output_shape,
+        axis=axis,
+        k=k,
+        largest=largest,
+        sorted=sorted_output,
+        input_dtype=input_dtype,
+        output_values_dtype=values_dtype,
+        output_indices_dtype=indices_dtype,
+    )
+
+
+register_lowering("TopK")(lower_topk)

emx_onnx_cgen/lowering/trilu.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import TriluOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import optional_name, value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+def _read_k_initializer(initializer: Initializer, node: Node) -> int:
+    if initializer.type.dtype != ScalarType.I64:
+        raise UnsupportedOpError(
+            f"{node.op_type} k input must be int64"
+        )
+    data = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if data.size != 1:
+        raise UnsupportedOpError(f"{node.op_type} k input must be scalar")
+    return int(data[0])
+
+
+@register_lowering("Trilu")
+def lower_trilu(graph: Graph, node: Node) -> TriluOp:
+    if len(node.inputs) not in {1, 2} or len(node.outputs) != 1:
+        raise UnsupportedOpError("Trilu must have 1 or 2 inputs and 1 output")
+    input_name = node.inputs[0]
+    if not input_name:
+        raise UnsupportedOpError("Trilu input must be provided")
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if input_shape != output_shape:
+        raise ShapeInferenceError("Trilu input and output shapes must match")
+    if len(output_shape) < 2:
+        raise UnsupportedOpError("Trilu expects input rank >= 2")
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Trilu expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    upper_attr = node.attrs.get("upper", 1)
+    upper = bool(int(upper_attr))
+    k_input = optional_name(node.inputs, 1)
+    k_value = 0
+    k_input_name = None
+    k_input_shape = None
+    k_input_dtype = None
+    if k_input:
+        k_initializer = _find_initializer(graph, k_input)
+        if k_initializer is not None:
+            k_value = _read_k_initializer(k_initializer, node)
+        else:
+            k_shape = value_shape(graph, k_input, node)
+            if not _is_scalar_shape(k_shape):
+                raise UnsupportedOpError("Trilu k input must be scalar")
+            k_dtype = value_dtype(graph, k_input, node)
+            if k_dtype != ScalarType.I64:
+                raise UnsupportedOpError("Trilu k input must be int64")
+            k_input_name = k_input
+            k_input_shape = k_shape
+            k_input_dtype = k_dtype
+    return TriluOp(
+        input0=input_name,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        upper=upper,
+        k_value=k_value,
+        k_input=k_input_name,
+        k_input_shape=k_input_shape,
+        k_input_dtype=k_input_dtype,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/unsqueeze.py

@@ -131,11 +131,11 @@ def lower_unsqueeze(graph: Graph, node: Node) -> ReshapeOp:
         for dim in output_shape:
             if input_index < len(input_shape) and dim == input_shape[input_index]:
                 input_index += 1
-                continue
-            if dim != 1:
-                raise ShapeInferenceError(
-                    "Unsqueeze output shape must insert ones only"
-                )
+            else:
+                if dim != 1:
+                    raise ShapeInferenceError(
+                        "Unsqueeze output shape must insert ones only"
+                    )
         if input_index != len(input_shape):
             raise ShapeInferenceError(
                 "Unsqueeze output shape must contain input shape in order"

emx_onnx_cgen/onnx_import.py

@@ -212,6 +212,9 @@ def import_onnx(model: onnx.ModelProto) -> Graph:
     dim_param_by_name = _collect_dim_params(
         tuple(model.graph.input) + tuple(model.graph.output)
     )
+    opset_imports = tuple(
+        (opset.domain, opset.version) for opset in model.opset_import
+    )
     try:
         model = shape_inference.infer_shapes(model, data_prop=True)
     except Exception as exc:  # pragma: no cover - onnx inference errors
@@ -258,4 +261,5 @@ def import_onnx(model: onnx.ModelProto) -> Graph:
         nodes=nodes,
         initializers=initializers,
         values=values,
+        opset_imports=opset_imports,
     )

emx_onnx_cgen/onnxruntime_utils.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def make_deterministic_session_options(ort: Any) -> Any:
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    return options

emx_onnx_cgen/ops.py

@@ -87,6 +87,7 @@ UNARY_OP_TYPES = {
     "Identity",
     "LeakyRelu",
     "Log",
+    "Mish",
     "Neg",
     "Not",
     "Reciprocal",
@@ -177,6 +178,7 @@ UNARY_SYMBOLS_DOUBLE = {
     ScalarFunction.LEAKY_RELU: "leaky_relu",
     ScalarFunction.POSITIVE: "identity",
     ScalarFunction.LOG: "log",
+    ScalarFunction.MISH: "mish",
     ScalarFunction.NEG: "neg",
     ScalarFunction.RECIPROCAL: "reciprocal",
     ScalarFunction.RELU: "relu",
@@ -215,6 +217,7 @@ UNARY_SYMBOLS_FLOAT = {
     ScalarFunction.LEAKY_RELU: "leaky_relu",
     ScalarFunction.POSITIVE: "identity",
     ScalarFunction.LOG: "logf",
+    ScalarFunction.MISH: "mish",
     ScalarFunction.NEG: "neg",
     ScalarFunction.RECIPROCAL: "reciprocal",
     ScalarFunction.RELU: "relu",
@@ -457,6 +460,7 @@ UNARY_APPLY_FUNCS = {
     "thresholded_relu": lambda value: np.where(
         value > 1.0, value, 0.0
     ),
+    "mish": lambda value: value * np.tanh(np.log1p(np.exp(value))),
     "atanhf": np.arctanh,
     "atanh": np.arctanh,
 }