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

emx_onnx_cgen/lowering/matmul.py

@@ -1,119 +1,17 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-
 from ..ir.ops import MatMulOp
-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 MatMulSpec:
-    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
-
-
-def resolve_matmul_spec(graph: Graph, node: Node) -> MatMulSpec:
-    if len(node.inputs) != 2 or len(node.outputs) != 1:
-        raise UnsupportedOpError("MatMul must have 2 inputs and 1 output")
-    input0_shape = _value_shape(graph, node.inputs[0], node)
-    input1_shape = _value_shape(graph, node.inputs[1], node)
-    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], node
-        )
-    )
-    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)
-    expected_output_shape = _value_shape(graph, node.outputs[0], node)
-    if expected_output_shape != output_shape:
-        raise ShapeInferenceError(
-            "MatMul output shape must be "
-            f"{output_shape}, got {expected_output_shape}"
-        )
-    return MatMulSpec(
-        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,
-    )
-
-
-def _broadcast_batch_shapes(
-    left: tuple[int, ...], right: tuple[int, ...], node: Node
-) -> 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 = []
-    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
-
-
 @register_lowering("MatMul")
 def lower_matmul(graph: Graph, node: Node) -> MatMulOp:
-    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
-    spec = resolve_matmul_spec(graph, node)
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("MatMul must have 2 inputs and 1 output")
     return MatMulOp(
         input0=node.inputs[0],
         input1=node.inputs[1],
         output=node.outputs[0],
-        input0_shape=spec.input0_shape,
-        input1_shape=spec.input1_shape,
-        output_shape=spec.output_shape,
-        batch_shape=spec.batch_shape,
-        input0_batch_shape=spec.input0_batch_shape,
-        input1_batch_shape=spec.input1_batch_shape,
-        m=spec.m,
-        n=spec.n,
-        k=spec.k,
-        left_vector=spec.left_vector,
-        right_vector=spec.right_vector,
-        dtype=op_dtype,
     )

emx_onnx_cgen/lowering/optional_has_element.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from ..errors import UnsupportedOpError
+from ..ir.context import GraphContext
+from ..ir.model import Node
+from ..ir.ops import OptionalHasElementOp
+from .registry import register_lowering
+
+
+@register_lowering("OptionalHasElement")
+def lower_optional_has_element(
+    ctx: GraphContext, node: Node
+) -> OptionalHasElementOp:
+    if len(node.inputs) != 1 or not node.inputs[0]:
+        raise UnsupportedOpError(
+            "OptionalHasElement expects exactly one non-empty input."
+        )
+    if len(node.outputs) != 1 or not node.outputs[0]:
+        raise UnsupportedOpError(
+            "OptionalHasElement expects exactly one output."
+        )
+    input_name = node.inputs[0]
+    value = ctx.find_value(input_name)
+    if not value.type.is_optional:
+        raise UnsupportedOpError(
+            "OptionalHasElement expects an optional input."
+        )
+    return OptionalHasElementOp(input0=input_name, output=node.outputs[0])

emx_onnx_cgen/lowering/qlinear_mul.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..ir.op_base import BroadcastingOpBase
+from ..ir.ops import QLinearMulOp
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+def _ensure_scalar_input(
+    graph: Graph, name: str, node: Node, label: str
+) -> tuple[int, ...]:
+    shape = _value_shape(graph, name, node)
+    if shape not in {(), (1,)}:
+        raise UnsupportedOpError(
+            f"QLinearMul {label} must be scalar, got shape {shape}"
+        )
+    return shape
+
+
+def _ensure_scale_dtype(dtype: ScalarType, label: str) -> None:
+    if not dtype.is_float:
+        raise UnsupportedOpError(
+            f"QLinearMul {label} must be float16/float/double"
+        )
+
+
+@register_lowering("QLinearMul")
+def lower_qlinear_mul(graph: Graph, node: Node) -> QLinearMulOp:
+    if len(node.inputs) != 8 or len(node.outputs) != 1:
+        raise UnsupportedOpError("QLinearMul must have 8 inputs and 1 output")
+    input0_shape = _value_shape(graph, node.inputs[0], node)
+    input1_shape = _value_shape(graph, node.inputs[3], node)
+    output_shape = BroadcastingOpBase.broadcast_shapes(
+        input0_shape, input1_shape
+    )
+    expected_output_shape = _value_shape(graph, node.outputs[0], node)
+    if expected_output_shape != output_shape:
+        raise ShapeInferenceError(
+            "QLinearMul output shape must be "
+            f"{output_shape}, got {expected_output_shape}"
+        )
+    input0_dtype = _value_dtype(graph, node.inputs[0], node)
+    input1_dtype = _value_dtype(graph, node.inputs[3], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if input0_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError("QLinearMul supports uint8/int8 inputs only")
+    if input1_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError("QLinearMul supports uint8/int8 inputs only")
+    if output_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError(
+            "QLinearMul supports uint8/int8 outputs only"
+        )
+    input0_scale_dtype = _value_dtype(graph, node.inputs[1], node)
+    input1_scale_dtype = _value_dtype(graph, node.inputs[4], node)
+    output_scale_dtype = _value_dtype(graph, node.inputs[6], node)
+    _ensure_scale_dtype(input0_scale_dtype, "a_scale")
+    _ensure_scale_dtype(input1_scale_dtype, "b_scale")
+    _ensure_scale_dtype(output_scale_dtype, "y_scale")
+    input0_zero_dtype = _value_dtype(graph, node.inputs[2], node)
+    input1_zero_dtype = _value_dtype(graph, node.inputs[5], node)
+    output_zero_dtype = _value_dtype(graph, node.inputs[7], node)
+    if input0_zero_dtype != input0_dtype:
+        raise UnsupportedOpError("QLinearMul a_zero_point dtype must match a")
+    if input1_zero_dtype != input1_dtype:
+        raise UnsupportedOpError("QLinearMul b_zero_point dtype must match b")
+    if output_zero_dtype != output_dtype:
+        raise UnsupportedOpError("QLinearMul y_zero_point dtype must match y")
+    input0_scale_shape = _ensure_scalar_input(
+        graph, node.inputs[1], node, "a_scale"
+    )
+    input1_scale_shape = _ensure_scalar_input(
+        graph, node.inputs[4], node, "b_scale"
+    )
+    output_scale_shape = _ensure_scalar_input(
+        graph, node.inputs[6], node, "y_scale"
+    )
+    input0_zero_shape = _ensure_scalar_input(
+        graph, node.inputs[2], node, "a_zero_point"
+    )
+    input1_zero_shape = _ensure_scalar_input(
+        graph, node.inputs[5], node, "b_zero_point"
+    )
+    output_zero_shape = _ensure_scalar_input(
+        graph, node.inputs[7], node, "y_zero_point"
+    )
+    return QLinearMulOp(
+        input0=node.inputs[0],
+        input0_scale=node.inputs[1],
+        input0_zero_point=node.inputs[2],
+        input1=node.inputs[3],
+        input1_scale=node.inputs[4],
+        input1_zero_point=node.inputs[5],
+        output_scale=node.inputs[6],
+        output_zero_point=node.inputs[7],
+        output=node.outputs[0],
+        input0_shape=input0_shape,
+        input1_shape=input1_shape,
+        output_shape=output_shape,
+        input0_dtype=input0_dtype,
+        input1_dtype=input1_dtype,
+        dtype=output_dtype,
+        input0_scale_dtype=input0_scale_dtype,
+        input1_scale_dtype=input1_scale_dtype,
+        output_scale_dtype=output_scale_dtype,
+        input0_scale_shape=input0_scale_shape,
+        input1_scale_shape=input1_scale_shape,
+        output_scale_shape=output_scale_shape,
+        input0_zero_shape=input0_zero_shape,
+        input1_zero_shape=input1_zero_shape,
+        output_zero_shape=output_zero_shape,
+    )

emx_onnx_cgen/lowering/reduce.py

@@ -525,17 +525,12 @@ def lower_reduce(graph: Graph, node: Node) -> ReduceOp | ReshapeOp:
     return ReduceOp(
         input0=node.inputs[0],
         output=node.outputs[0],
-        input_shape=input_shape,
-        output_shape=spec.output_shape,
         axes=spec.axes or (),
         axes_input=spec.axes_input,
-        axes_input_shape=spec.axes_input_shape,
-        axes_input_dtype=spec.axes_input_dtype,
         keepdims=spec.keepdims,
         noop_with_empty_axes=bool(int(node.attrs.get("noop_with_empty_axes", 0))),
         reduce_kind=REDUCE_KIND_BY_OP[node.op_type],
         reduce_count=spec.reduce_count,
-        dtype=op_dtype,
     )
 
 

emx_onnx_cgen/lowering/reshape.py

@@ -2,31 +2,20 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Initializer, Node
-from .common import value_shape as resolved_value_shape
+from ..ir.ops import ReshapeOp
+from .common import value_dtype, value_shape as resolved_value_shape
 from .registry import register_lowering
 
 
 def _value_shape(graph: Graph, name: str, node: Node) -> tuple[int, ...]:
-    try:
-        return graph.find_value(name).type.shape
-    except KeyError as exc:
-        raise ShapeInferenceError(
-            f"Missing shape for value '{name}' in op {node.op_type}. "
-            "Hint: run ONNX shape inference or export with static shapes."
-        ) from exc
+    return resolved_value_shape(graph, name, node)
 
 
 def _value_dtype(graph: Graph, name: str, node: Node) -> ScalarType:
-    try:
-        return graph.find_value(name).type.dtype
-    except KeyError as exc:
-        raise ShapeInferenceError(
-            f"Missing dtype for value '{name}' in op {node.op_type}. "
-            "Hint: run ONNX shape inference or export with static shapes."
-        ) from exc
+    return value_dtype(graph, name, node)
 
 
 def _shape_product(shape: tuple[int, ...]) -> int:
@@ -350,6 +339,8 @@ def lower_reshape(graph: Graph, node: Node) -> ReshapeOp:
     for dim in output_shape:
         if dim < 0:
             raise ShapeInferenceError("Dynamic dims are not supported")
+    if isinstance(graph, GraphContext):
+        graph.set_shape(node.outputs[0], output_shape)
     return ReshapeOp(
         input0=node.inputs[0],
         output=node.outputs[0],

emx_onnx_cgen/lowering/shape.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..ir.ops import ShapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
-from .common import value_dtype, value_shape
+from ..ir.ops import ShapeOp
+from .common import value_dtype, value_has_dim_params, value_shape
 from .registry import register_lowering
 
 
@@ -29,10 +30,13 @@ def lower_shape(graph: Graph, node: Node) -> ShapeOp:
         raise UnsupportedOpError("Shape must have 1 input and 1 output")
     input_shape = value_shape(graph, node.inputs[0], node)
     output_shape = value_shape(graph, node.outputs[0], node)
-    if len(output_shape) != 1:
-        raise ShapeInferenceError("Shape output must be 1D")
-    if output_shape[0] < 0:
-        raise ShapeInferenceError("Shape output length must be non-negative")
+    if value_has_dim_params(graph, node.outputs[0]) or not output_shape:
+        output_shape = ()
+    if output_shape:
+        if len(output_shape) != 1:
+            raise ShapeInferenceError("Shape output must be 1D")
+        if output_shape[0] < 0:
+            raise ShapeInferenceError("Shape output length must be non-negative")
     input_dtype = value_dtype(graph, node.inputs[0], node)
     output_dtype = value_dtype(graph, node.outputs[0], node)
     if output_dtype != ScalarType.I64:
@@ -43,16 +47,18 @@ def lower_shape(graph: Graph, node: Node) -> ShapeOp:
         len(input_shape), start=start, end=end
     )
     expected_shape = (max(0, end_index - start_index),)
-    if expected_shape != output_shape:
+    if output_shape and expected_shape != output_shape:
         raise ShapeInferenceError(
             "Shape output shape must be "
             f"{expected_shape}, got {output_shape}"
         )
+    if isinstance(graph, GraphContext):
+        graph.set_shape(node.outputs[0], expected_shape)
     return ShapeOp(
         input0=node.inputs[0],
         output=node.outputs[0],
         input_shape=input_shape,
-        output_shape=output_shape,
+        output_shape=expected_shape,
         values=input_shape[start_index:end_index],
         dtype=output_dtype,
         input_dtype=input_dtype,

emx_onnx_cgen/lowering/slice.py

@@ -6,10 +6,16 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..ir.ops import SliceOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Initializer, Node
-from ..lowering.common import value_dtype, value_shape
+from ..ir.ops import SliceOp
+from ..lowering.common import (
+    resolve_int_list_from_value,
+    value_has_dim_params,
+    value_dtype,
+    value_shape,
+)
 from ..validation import normalize_axis
 from .registry import register_lowering
 
@@ -70,7 +76,7 @@ def _maybe_read_int_list(
 ) -> list[int] | None:
     initializer = _find_initializer(graph, name)
     if initializer is None:
-        return None
+        return resolve_int_list_from_value(graph, name, node)
     return _read_int_list(graph, name, node, label=label)
 
 
@@ -335,6 +341,8 @@ def resolve_slice_spec(graph: Graph, node: Node) -> SliceSpec:
 def lower_slice(graph: Graph, node: Node) -> SliceOp:
     input_shape = value_shape(graph, node.inputs[0], node)
     output_shape = value_shape(graph, node.outputs[0], node)
+    if value_has_dim_params(graph, node.outputs[0]):
+        output_shape = ()
     input_dtype = value_dtype(graph, node.inputs[0], node)
     output_dtype = value_dtype(graph, node.outputs[0], node)
     if input_dtype != output_dtype:
@@ -356,6 +364,8 @@ def lower_slice(graph: Graph, node: Node) -> SliceOp:
                 f"{node.op_type} output shape must be "
                 f"{computed_output_shape}, got {output_shape}"
             )
+        if isinstance(graph, GraphContext):
+            graph.set_shape(node.outputs[0], computed_output_shape)
         return SliceOp(
             input0=node.inputs[0],
             output=node.outputs[0],
@@ -379,7 +389,7 @@ def lower_slice(graph: Graph, node: Node) -> SliceOp:
             dtype=input_dtype,
             input_dtype=input_dtype,
         )
-    if len(output_shape) != len(input_shape):
+    if output_shape and len(output_shape) != len(input_shape):
         raise ShapeInferenceError(
             f"{node.op_type} output rank must match input rank"
         )

emx_onnx_cgen/lowering/softmax.py

@@ -3,49 +3,15 @@ from __future__ import annotations
 from ..ir.ops 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
-from ..validation import ensure_output_shape_matches_input
-from ..validation import normalize_axis as _normalize_axis
 
 
 @register_lowering("Softmax")
 def lower_softmax(graph: Graph, node: Node) -> SoftmaxOp:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
         raise UnsupportedOpError("Softmax must have 1 input and 1 output")
-    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
-    if not op_dtype.is_float:
-        raise UnsupportedOpError(
-            "Softmax supports float16, float, and double inputs only"
-        )
-    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(axis_attr),
-        input_shape,
-        node,
-    )
-    outer = _shape_product(input_shape[:axis]) if axis > 0 else 1
-    axis_size = input_shape[axis]
-    inner = (
-        _shape_product(input_shape[axis + 1 :])
-        if axis + 1 < len(input_shape)
-        else 1
-    )
     return SoftmaxOp(
         input0=node.inputs[0],
         output=node.outputs[0],
-        outer=outer,
-        axis_size=axis_size,
-        inner=inner,
-        axis=axis,
-        shape=input_shape,
-        dtype=op_dtype,
+        axis=int(node.attrs["axis"]) if "axis" in node.attrs else None,
     )

emx_onnx_cgen/lowering/split.py

@@ -6,8 +6,14 @@ from shared.scalar_types import ScalarType
 
 from ..ir.ops import SplitOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Initializer, Node
-from ..lowering.common import optional_name, value_dtype, value_shape
+from ..lowering.common import (
+    optional_name,
+    resolve_int_list_from_value,
+    value_dtype,
+    value_shape,
+)
 from ..validation import normalize_axis
 from .registry import register_lowering
 
@@ -46,6 +52,22 @@ def _validate_static_dims(shape: tuple[int, ...], node: Node) -> None:
         )
 
 
+def _validate_output_ranks(
+    output_shapes: list[tuple[int, ...]],
+    input_shape: tuple[int, ...],
+    node: Node,
+) -> None:
+    expected_rank = len(input_shape)
+    for output_shape in output_shapes:
+        if not output_shape:
+            continue
+        if len(output_shape) != expected_rank:
+            raise ShapeInferenceError(
+                f"{node.op_type} output rank must match input rank "
+                f"{expected_rank}, got {len(output_shape)}"
+            )
+
+
 def _normalize_num_outputs(node: Node, output_count: int) -> int:
     num_outputs_attr = node.attrs.get("num_outputs")
     if num_outputs_attr is None:
@@ -75,6 +97,7 @@ def lower_split(graph: Graph, node: Node) -> SplitOp:
     output_shapes = [
         value_shape(graph, output, node) for output in node.outputs
     ]
+    _validate_output_ranks(output_shapes, input_shape, node)
     input_dtype = value_dtype(graph, input_name, node)
     output_dtypes = {value_dtype(graph, output, node) for output in node.outputs}
     if output_dtypes != {input_dtype}:
@@ -107,7 +130,15 @@ def lower_split(graph: Graph, node: Node) -> SplitOp:
                 raise ShapeInferenceError(
                     f"Split expects {len(node.outputs)} outputs, got {split_shape[0]}"
                 )
-            split_sizes = [shape[axis] for shape in output_shapes]
+            split_sizes = resolve_int_list_from_value(graph, split_name, node)
+            if split_sizes is None:
+                if all(output_shape for output_shape in output_shapes):
+                    split_sizes = [shape[axis] for shape in output_shapes]
+                else:
+                    raise ShapeInferenceError(
+                        "Split sizes must be constant when output shapes "
+                        "are unavailable"
+                    )
         if len(split_sizes) != len(node.outputs):
             raise ShapeInferenceError(
                 f"Split expects {len(split_sizes)} outputs, got {len(node.outputs)}"
@@ -133,11 +164,14 @@ def lower_split(graph: Graph, node: Node) -> SplitOp:
         shape = list(input_shape)
         shape[axis] = size
         computed_shape = tuple(shape)
-        if output_shape != computed_shape:
+        if output_shape and output_shape != computed_shape:
             raise ShapeInferenceError(
                 f"Split output shape must be {computed_shape}, got {output_shape}"
             )
         computed_shapes.append(computed_shape)
+    if isinstance(graph, GraphContext):
+        for output_name, shape in zip(node.outputs, computed_shapes):
+            graph.set_shape(output_name, shape)
     return SplitOp(
         input0=input_name,
         outputs=tuple(node.outputs),