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

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

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import NonZeroOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+@register_lowering("NonZero")
+def lower_nonzero(graph: Graph, node: Node) -> NonZeroOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("NonZero must have 1 input and 1 output")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    if len(input_shape) == 0:
+        raise UnsupportedOpError("NonZero does not support scalar inputs")
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if len(output_shape) != 2:
+        raise ShapeInferenceError("NonZero output must be 2D")
+    if output_shape[0] != len(input_shape):
+        raise ShapeInferenceError(
+            "NonZero output shape must be "
+            f"({len(input_shape)}, N), got {output_shape}"
+        )
+    if output_shape[0] < 0 or output_shape[1] < 0:
+        raise ShapeInferenceError(
+            "NonZero output shape must be non-negative"
+        )
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if output_dtype != ScalarType.I64:
+        raise UnsupportedOpError("NonZero output dtype must be int64")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    return NonZeroOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/one_hot.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+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
+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_scalar_initializer(
+    graph: Graph, name: str, node: Node, label: str
+) -> int | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    data = np.array(initializer.data)
+    if data.size != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} {label} input must be a scalar"
+        )
+    return int(data.reshape(-1)[0])
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+def _normalize_onehot_axis(axis: int, rank: int, node: Node) -> int:
+    if axis < 0:
+        axis += rank + 1
+    if axis < 0 or axis > rank:
+        raise ShapeInferenceError(
+            f"{node.op_type} axis {axis} is out of range for rank {rank}"
+        )
+    return axis
+
+
+@register_lowering("OneHot")
+def lower_onehot(graph: Graph, node: Node) -> OneHotOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError("OneHot must have 3 inputs and 1 output")
+    indices_name, depth_name, values_name = node.inputs
+    indices_shape = value_shape(graph, indices_name, node)
+    depth_shape = value_shape(graph, depth_name, node)
+    values_shape = value_shape(graph, values_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if not _is_scalar_shape(depth_shape):
+        raise UnsupportedOpError("OneHot depth input must be a scalar")
+    if len(values_shape) != 1 or values_shape[0] != 2:
+        raise UnsupportedOpError(
+            "OneHot values input must be a 1D tensor of size 2"
+        )
+    output_rank = len(indices_shape) + 1
+    if len(output_shape) != output_rank:
+        raise ShapeInferenceError(
+            f"OneHot output rank must be {output_rank}, got {len(output_shape)}"
+        )
+    axis = _normalize_onehot_axis(
+        int(node.attrs.get("axis", -1)), len(indices_shape), node
+    )
+    depth_value = _read_scalar_initializer(graph, depth_name, node, "depth")
+    if depth_value is not None:
+        if depth_value < 0:
+            raise ShapeInferenceError("OneHot depth must be non-negative")
+        if output_shape[axis] != depth_value:
+            raise ShapeInferenceError(
+                "OneHot output depth must be "
+                f"{depth_value}, got {output_shape[axis]}"
+            )
+        depth_dim = depth_value
+    else:
+        depth_dim = output_shape[axis]
+        if depth_dim < 0:
+            raise ShapeInferenceError("OneHot output depth must be non-negative")
+    expected_output_shape = (
+        indices_shape[:axis] + (depth_dim,) + indices_shape[axis:]
+    )
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "OneHot output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    indices_dtype = value_dtype(graph, indices_name, node)
+    depth_dtype = value_dtype(graph, depth_name, node)
+    values_dtype = value_dtype(graph, values_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if indices_dtype.is_bool:
+        raise UnsupportedOpError("OneHot indices must be numeric")
+    if depth_dtype.is_bool:
+        raise UnsupportedOpError("OneHot depth must be numeric")
+    if values_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "OneHot values dtype must match output dtype, "
+            f"got {values_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    return OneHotOp(
+        indices=indices_name,
+        depth=depth_name,
+        values=values_name,
+        output=node.outputs[0],
+        axis=axis,
+        indices_shape=indices_shape,
+        values_shape=values_shape,
+        output_shape=output_shape,
+        depth_dim=depth_dim,
+        dtype=values_dtype,
+        indices_dtype=indices_dtype,
+        depth_dtype=depth_dtype,
+    )

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

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..dtypes import scalar_type_from_onnx
+from ..errors import ShapeInferenceError, UnsupportedOpError
+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 ..ir.ops import QuantizeLinearOp
+
+
+@dataclass(frozen=True)
+class QuantizeSpec:
+    input_shape: tuple[int, ...]
+    scale_shape: tuple[int, ...]
+    axis: int | None
+    output_dtype: ScalarType
+
+
+def _resolve_output_dtype(
+    graph: Graph, node: Node, zero_point_name: str | None
+) -> ScalarType:
+    output_attr = int(node.attrs.get("output_dtype", 0))
+    if output_attr:
+        output_dtype = _value_dtype(graph, node.outputs[0], node)
+        attr_dtype = scalar_type_from_onnx(output_attr)
+        if attr_dtype is None:
+            raise UnsupportedOpError(
+                "QuantizeLinear output_dtype must map to a supported scalar type"
+            )
+        if output_dtype != attr_dtype:
+            raise UnsupportedOpError(
+                "QuantizeLinear output_dtype must match output tensor dtype"
+            )
+        return output_dtype
+    if zero_point_name is None:
+        return ScalarType.U8
+    return _value_dtype(graph, zero_point_name, node)
+
+
+def resolve_quantize_spec(graph: Graph, node: Node) -> QuantizeSpec:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "QuantizeLinear must have 2 or 3 inputs and 1 output"
+        )
+    supported_attrs = {"axis", "block_size", "output_dtype", "precision", "saturate"}
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("QuantizeLinear has unsupported attributes")
+    block_size = int(node.attrs.get("block_size", 0))
+    if block_size != 0:
+        raise UnsupportedOpError("QuantizeLinear block_size is not supported")
+    precision = int(node.attrs.get("precision", 0))
+    if precision != 0:
+        raise UnsupportedOpError("QuantizeLinear precision is not supported")
+    saturate = int(node.attrs.get("saturate", 1))
+    if saturate != 1:
+        raise UnsupportedOpError("QuantizeLinear saturate must be 1")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    scale_shape = _value_shape(graph, node.inputs[1], node)
+    zero_point_name = optional_name(node.inputs, 2)
+    output_dtype = _resolve_output_dtype(graph, node, zero_point_name)
+    if output_dtype not in {
+        ScalarType.U8,
+        ScalarType.I8,
+        ScalarType.U16,
+        ScalarType.I16,
+    }:
+        raise UnsupportedOpError(
+            "QuantizeLinear supports int8/uint8/int16/uint16 outputs only"
+        )
+    if zero_point_name is not None:
+        zero_point_dtype = _value_dtype(graph, zero_point_name, node)
+        if zero_point_dtype != output_dtype:
+            raise UnsupportedOpError(
+                "QuantizeLinear zero_point dtype must match output dtype"
+            )
+        zero_point_shape = _value_shape(graph, zero_point_name, node)
+        if zero_point_shape != scale_shape:
+            raise ShapeInferenceError(
+                "QuantizeLinear zero_point shape must match scale shape"
+            )
+    if scale_shape not in {(), (1,)}:
+        if len(scale_shape) != 1:
+            raise UnsupportedOpError(
+                "QuantizeLinear supports per-tensor and per-axis scales only"
+            )
+        axis = int(node.attrs.get("axis", 1))
+        axis = normalize_axis(axis, input_shape, node)
+        if scale_shape[0] != input_shape[axis]:
+            raise ShapeInferenceError(
+                "QuantizeLinear scale length must match input axis size"
+            )
+    else:
+        axis = None
+    return QuantizeSpec(
+        input_shape=input_shape,
+        scale_shape=scale_shape,
+        axis=axis,
+        output_dtype=output_dtype,
+    )
+
+
+@register_lowering("QuantizeLinear")
+def lower_quantize_linear(graph: Graph, node: Node) -> QuantizeLinearOp:
+    op_dtype = _value_dtype(graph, node.inputs[0], node)
+    scale_dtype = _value_dtype(graph, node.inputs[1], node)
+    if not op_dtype.is_float or not scale_dtype.is_float:
+        raise UnsupportedOpError(
+            "QuantizeLinear supports float16/float/double inputs only"
+        )
+    spec = resolve_quantize_spec(graph, node)
+    return QuantizeLinearOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        zero_point=optional_name(node.inputs, 2),
+        output=node.outputs[0],
+        input_shape=spec.input_shape,
+        axis=spec.axis,
+        dtype=spec.output_dtype,
+        input_dtype=op_dtype,
+        scale_dtype=scale_dtype,
+    )

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