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

emx_onnx_cgen/lowering/gather_elements.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import GatherElementsOp
+from ..ir.ops import GatherElementsOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import normalize_axis

emx_onnx_cgen/lowering/gather_nd.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import GatherNDOp
+from ..ir.ops import GatherNDOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import value_dtype as _value_dtype

emx_onnx_cgen/lowering/gemm.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import GemmOp
+from ..ir.ops import GemmOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/global_max_pool.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReduceOp
+from ..ir.ops import ReduceOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import value_dtype as _value_dtype

emx_onnx_cgen/lowering/grid_sample.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import GridSampleOp
+from ..ir.ops import GridSampleOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import value_dtype, value_shape

emx_onnx_cgen/lowering/group_normalization.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import GroupNormalizationOp
+from ..ir.ops import GroupNormalizationOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import ensure_output_shape_matches_input

emx_onnx_cgen/lowering/hardmax.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import HardmaxOp
+from ..ir.ops import HardmaxOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/identity.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import IdentityOp
+from ..ir.ops import IdentityOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import value_dtype, value_shape

emx_onnx_cgen/lowering/instance_normalization.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import InstanceNormalizationOp
+from ..ir.ops import InstanceNormalizationOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import ensure_output_shape_matches_input

emx_onnx_cgen/lowering/layer_normalization.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import LayerNormalizationOp
+from ..ir.ops import LayerNormalizationOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import ensure_output_shape_matches_input

emx_onnx_cgen/lowering/logsoftmax.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import LogSoftmaxOp
+from ..ir.ops import LogSoftmaxOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/lp_normalization.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import LpNormalizationOp
+from ..ir.ops import LpNormalizationOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import ensure_output_shape_matches_input

emx_onnx_cgen/lowering/lp_pool.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from ..codegen.c_emitter import LpPoolOp
+from ..ir.ops import LpPoolOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/lrn.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from ..codegen.c_emitter import LrnOp
+from ..ir.ops import LrnOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/lstm.py

@@ -323,7 +323,7 @@ def resolve_lstm_spec(graph: Graph, node: Node) -> LstmSpec:
 
 @register_lowering("LSTM")
 def lower_lstm(graph: Graph, node: Node) -> "LstmOp":
-    from ..codegen.c_emitter import LstmOp
+    from ..ir.ops import LstmOp
 
     spec = resolve_lstm_spec(graph, node)
     return LstmOp(

emx_onnx_cgen/lowering/matmul.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from ..codegen.c_emitter import MatMulOp
+from ..ir.ops import MatMulOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/maxpool.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import MaxPoolOp
+from ..ir.ops import MaxPoolOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/mean_variance_normalization.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import MeanVarianceNormalizationOp
+from ..ir.ops import MeanVarianceNormalizationOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import ensure_output_shape_matches_input

emx_onnx_cgen/lowering/negative_log_likelihood_loss.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import NegativeLogLikelihoodLossOp
+from ..ir.ops import NegativeLogLikelihoodLossOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .common import shape_product as _shape_product

emx_onnx_cgen/lowering/non_max_suppression.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import NonMaxSuppressionOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..lowering.common import optional_name, shape_product, value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _validate_scalar_input(
+    graph: Graph,
+    name: str,
+    node: Node,
+    *,
+    allowed_dtypes: set[ScalarType],
+    label: str,
+) -> tuple[ScalarType, tuple[int, ...]]:
+    dtype = value_dtype(graph, name, node)
+    if dtype not in allowed_dtypes:
+        allowed = ", ".join(sorted(d.onnx_name for d in allowed_dtypes))
+        raise UnsupportedOpError(
+            f"{node.op_type} {label} must be {allowed}, got {dtype.onnx_name}"
+        )
+    shape = value_shape(graph, name, node)
+    if shape not in {(), (1,)}:
+        total = shape_product(shape)
+        if total != 1:
+            raise ShapeInferenceError(
+                f"{node.op_type} {label} must be a scalar tensor, got shape {shape}"
+            )
+    return dtype, shape
+
+
+@register_lowering("NonMaxSuppression")
+def lower_non_max_suppression(graph: Graph, node: Node) -> NonMaxSuppressionOp:
+    if node.op_type != "NonMaxSuppression":
+        raise UnsupportedOpError(f"Unsupported op {node.op_type}")
+    if len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} must have 1 output, got {len(node.outputs)}"
+        )
+    if len(node.inputs) < 2 or len(node.inputs) > 5:
+        raise UnsupportedOpError(
+            f"{node.op_type} must have 2 to 5 inputs, got {len(node.inputs)}"
+        )
+
+    boxes = node.inputs[0]
+    scores = node.inputs[1]
+    max_output_boxes_per_class = optional_name(node.inputs, 2)
+    iou_threshold = optional_name(node.inputs, 3)
+    score_threshold = optional_name(node.inputs, 4)
+    output = node.outputs[0]
+
+    boxes_shape = value_shape(graph, boxes, node)
+    scores_shape = value_shape(graph, scores, node)
+    if len(boxes_shape) != 3 or boxes_shape[2] != 4:
+        raise ShapeInferenceError(
+            f"{node.op_type} boxes input must have shape "
+            f"[num_batches, num_boxes, 4], got {boxes_shape}"
+        )
+    if len(scores_shape) != 3:
+        raise ShapeInferenceError(
+            f"{node.op_type} scores input must have shape "
+            f"[num_batches, num_classes, num_boxes], got {scores_shape}"
+        )
+    if boxes_shape[0] != scores_shape[0]:
+        raise ShapeInferenceError(
+            f"{node.op_type} boxes/scores batch dims must match, "
+            f"got {boxes_shape[0]} and {scores_shape[0]}"
+        )
+    if boxes_shape[1] != scores_shape[2]:
+        raise ShapeInferenceError(
+            f"{node.op_type} boxes num_boxes dim {boxes_shape[1]} "
+            f"must match scores num_boxes dim {scores_shape[2]}"
+        )
+
+    boxes_dtype = value_dtype(graph, boxes, node)
+    scores_dtype = value_dtype(graph, scores, node)
+    if boxes_dtype != scores_dtype or not boxes_dtype.is_float:
+        raise UnsupportedOpError(
+            f"{node.op_type} boxes and scores must be the same float dtype, "
+            f"got {boxes_dtype.onnx_name} and {scores_dtype.onnx_name}"
+        )
+
+    max_output_dtype = None
+    max_output_shape = None
+    if max_output_boxes_per_class is not None:
+        max_output_dtype, max_output_shape = _validate_scalar_input(
+            graph,
+            max_output_boxes_per_class,
+            node,
+            allowed_dtypes={ScalarType.I32, ScalarType.I64},
+            label="max_output_boxes_per_class input",
+        )
+
+    iou_threshold_dtype = None
+    iou_threshold_shape = None
+    if iou_threshold is not None:
+        iou_threshold_dtype, iou_threshold_shape = _validate_scalar_input(
+            graph,
+            iou_threshold,
+            node,
+            allowed_dtypes={ScalarType.F32, ScalarType.F64},
+            label="iou_threshold input",
+        )
+
+    score_threshold_dtype = None
+    score_threshold_shape = None
+    if score_threshold is not None:
+        score_threshold_dtype, score_threshold_shape = _validate_scalar_input(
+            graph,
+            score_threshold,
+            node,
+            allowed_dtypes={ScalarType.F32, ScalarType.F64},
+            label="score_threshold input",
+        )
+
+    output_shape = value_shape(graph, output, node)
+    if len(output_shape) != 2 or output_shape[1] != 3:
+        raise ShapeInferenceError(
+            f"{node.op_type} output must have shape [num_selected, 3], "
+            f"got {output_shape}"
+        )
+    output_dtype = value_dtype(graph, output, node)
+    if output_dtype != ScalarType.I64:
+        raise UnsupportedOpError(
+            f"{node.op_type} output dtype must be int64"
+        )
+
+    center_point_box = int(node.attrs.get("center_point_box", 0))
+    if center_point_box not in {0, 1}:
+        raise UnsupportedOpError(
+            f"{node.op_type} center_point_box must be 0 or 1, got {center_point_box}"
+        )
+
+    return NonMaxSuppressionOp(
+        boxes=boxes,
+        scores=scores,
+        max_output_boxes_per_class=max_output_boxes_per_class,
+        iou_threshold=iou_threshold,
+        score_threshold=score_threshold,
+        output=output,
+        boxes_shape=boxes_shape,
+        scores_shape=scores_shape,
+        output_shape=output_shape,
+        center_point_box=center_point_box,
+        boxes_dtype=boxes_dtype,
+        output_dtype=output_dtype,
+        max_output_dtype=max_output_dtype,
+        max_output_shape=max_output_shape,
+        iou_threshold_dtype=iou_threshold_dtype,
+        iou_threshold_shape=iou_threshold_shape,
+        score_threshold_dtype=score_threshold_dtype,
+        score_threshold_shape=score_threshold_shape,
+    )

emx_onnx_cgen/lowering/nonzero.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import NonZeroOp
+from ..ir.ops import NonZeroOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import value_dtype, value_shape

emx_onnx_cgen/lowering/one_hot.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import OneHotOp
+from ..ir.ops import OneHotOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import value_dtype, value_shape

emx_onnx_cgen/lowering/pad.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import PadOp
+from ..ir.ops import PadOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import optional_name, value_dtype, value_shape

emx_onnx_cgen/lowering/qlinear_matmul.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import QLinearMatMulOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class QLinearMatMulSpec:
+    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_qlinear_matmul_spec(graph: Graph, node: Node) -> QLinearMatMulSpec:
+    if len(node.inputs) != 8 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "QLinearMatMul 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)
+    if len(input0_shape) < 1 or len(input1_shape) < 1:
+        raise UnsupportedOpError(
+            "QLinearMatMul 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(
+            "QLinearMatMul inner dimensions must match, "
+            f"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(
+            "QLinearMatMul output shape must be "
+            f"{output_shape}, got {expected_output_shape}"
+        )
+    return QLinearMatMulSpec(
+        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(
+                "QLinearMatMul 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 _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"QLinearMatMul {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"QLinearMatMul {label} must be float16/float/double"
+        )
+
+
+@register_lowering("QLinearMatMul")
+def lower_qlinear_matmul(graph: Graph, node: Node) -> QLinearMatMulOp:
+    spec = resolve_qlinear_matmul_spec(graph, node)
+    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(
+            "QLinearMatMul supports uint8/int8 inputs only"
+        )
+    if input1_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError(
+            "QLinearMatMul supports uint8/int8 inputs only"
+        )
+    if output_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError(
+            "QLinearMatMul 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(
+            "QLinearMatMul a_zero_point dtype must match a"
+        )
+    if input1_zero_dtype != input1_dtype:
+        raise UnsupportedOpError(
+            "QLinearMatMul b_zero_point dtype must match b"
+        )
+    if output_zero_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "QLinearMatMul 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 QLinearMatMulOp(
+        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=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,
+        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/quantize_linear.py

@@ -10,7 +10,7 @@ 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 ..codegen.c_emitter import QuantizeLinearOp
+from ..ir.ops import QuantizeLinearOp
 
 
 @dataclass(frozen=True)

emx_onnx_cgen/lowering/range.py

@@ -6,7 +6,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import RangeOp
+from ..ir.ops import RangeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import node_dtype, value_shape

emx_onnx_cgen/lowering/reduce.py

@@ -6,7 +6,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReduceOp, ReshapeOp
+from ..ir.ops import ReduceOp, ReshapeOp
 from ..dtypes import scalar_type_from_onnx
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node

emx_onnx_cgen/lowering/registry.py

@@ -3,32 +3,51 @@ from __future__ import annotations
 from collections.abc import Callable, Mapping
 from typing import TypeVar
 
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
+from ..ir.op_base import OpBase
 from ..errors import UnsupportedOpError
 
 LoweredOp = TypeVar("LoweredOp")
 Handler = TypeVar("Handler")
 
-_LOWERING_REGISTRY: dict[str, Callable[[Graph, Node], object]] = {}
+_LOWERING_REGISTRY: dict[str, Callable[[Graph | GraphContext, Node], OpBase]] = {}
 
 
 def register_lowering(
     op_type: str,
 ) -> Callable[[Callable[[Graph, Node], LoweredOp]], Callable[[Graph, Node], LoweredOp]]:
     def decorator(
-        func: Callable[[Graph, Node], LoweredOp],
-    ) -> Callable[[Graph, Node], LoweredOp]:
+        func: Callable[[Graph | GraphContext, Node], LoweredOp],
+    ) -> Callable[[Graph | GraphContext, Node], LoweredOp]:
         _LOWERING_REGISTRY[op_type] = func
         return func
 
     return decorator
 
 
-def get_lowering(op_type: str) -> Callable[[Graph, Node], object] | None:
+def register_lowering_if_missing(
+    op_type: str,
+) -> Callable[[Callable[[Graph | GraphContext, Node], LoweredOp]], Callable[[Graph | GraphContext, Node], LoweredOp]]:
+    def decorator(
+        func: Callable[[Graph | GraphContext, Node], LoweredOp],
+    ) -> Callable[[Graph | GraphContext, Node], LoweredOp]:
+        if op_type not in _LOWERING_REGISTRY:
+            _LOWERING_REGISTRY[op_type] = func
+        return func
+
+    return decorator
+
+
+def get_lowering(
+    op_type: str,
+) -> Callable[[Graph | GraphContext, Node], OpBase] | None:
     return _LOWERING_REGISTRY.get(op_type)
 
 
-def get_lowering_registry() -> Mapping[str, Callable[[Graph, Node], object]]:
+def get_lowering_registry() -> Mapping[
+    str, Callable[[Graph | GraphContext, Node], OpBase]
+]:
     return _LOWERING_REGISTRY