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

emx_onnx_cgen/lowering/average_pool.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from dataclasses import dataclass
 
 from ..ir.ops import AveragePoolOp
@@ -12,16 +13,26 @@ from .registry import register_lowering
 class _AveragePoolSpec:
     batch: int
     channels: int
+    spatial_rank: int
+    in_d: int
     in_h: int
     in_w: int
+    out_d: int
     out_h: int
     out_w: int
+    kernel_d: int
     kernel_h: int
     kernel_w: int
+    dilation_d: int
+    dilation_h: int
+    dilation_w: int
+    stride_d: int
     stride_h: int
     stride_w: int
+    pad_front: int
     pad_top: int
     pad_left: int
+    pad_back: int
     pad_bottom: int
     pad_right: int
     count_include_pad: bool
@@ -54,6 +65,7 @@ def _resolve_average_pool_spec(graph: Graph, node: Node) -> _AveragePoolSpec:
         "auto_pad",
         "ceil_mode",
         "count_include_pad",
+        "dilations",
         "kernel_shape",
         "pads",
         "strides",
@@ -63,11 +75,9 @@ def _resolve_average_pool_spec(graph: Graph, node: Node) -> _AveragePoolSpec:
     auto_pad = node.attrs.get("auto_pad", b"NOTSET")
     if isinstance(auto_pad, bytes):
         auto_pad = auto_pad.decode("utf-8", errors="ignore")
-    if auto_pad not in ("", "NOTSET"):
-        raise UnsupportedOpError("AveragePool supports auto_pad=NOTSET only")
     ceil_mode = int(node.attrs.get("ceil_mode", 0))
-    if ceil_mode != 0:
-        raise UnsupportedOpError("AveragePool supports ceil_mode=0 only")
+    if ceil_mode not in (0, 1):
+        raise UnsupportedOpError("AveragePool supports ceil_mode=0 or 1 only")
     count_include_pad = int(node.attrs.get("count_include_pad", 0))
     if count_include_pad not in (0, 1):
         raise UnsupportedOpError("AveragePool supports count_include_pad 0 or 1")
@@ -75,47 +85,128 @@ def _resolve_average_pool_spec(graph: Graph, node: Node) -> _AveragePoolSpec:
     if kernel_shape is None:
         raise UnsupportedOpError("AveragePool requires kernel_shape")
     kernel_shape = tuple(int(value) for value in kernel_shape)
-    if len(kernel_shape) != 2:
-        raise UnsupportedOpError("AveragePool expects 2D kernel_shape")
-    kernel_h, kernel_w = kernel_shape
-    strides = tuple(int(value) for value in node.attrs.get("strides", (1, 1)))
-    if len(strides) != 2:
-        raise UnsupportedOpError("AveragePool expects 2D strides")
-    pads = tuple(int(value) for value in node.attrs.get("pads", (0, 0, 0, 0)))
-    if len(pads) != 4:
-        raise UnsupportedOpError("AveragePool expects 4D pads")
-    pad_top, pad_left, pad_bottom, pad_right = pads
     input_shape = _value_shape(graph, node.inputs[0], node)
-    if len(input_shape) != 4:
-        raise UnsupportedOpError("AveragePool supports NCHW 2D inputs only")
-    batch, channels, in_h, in_w = input_shape
-    stride_h, stride_w = strides
-    out_h = (in_h + pad_top + pad_bottom - kernel_h) // stride_h + 1
-    out_w = (in_w + pad_left + pad_right - kernel_w) // stride_w + 1
-    if out_h < 0 or out_w < 0:
+    if len(input_shape) < 3:
+        raise UnsupportedOpError("AveragePool expects NCHW inputs with spatial dims")
+    spatial_rank = len(input_shape) - 2
+    if spatial_rank not in {1, 2, 3}:
+        raise UnsupportedOpError("AveragePool supports 1D/2D/3D inputs only")
+    if len(kernel_shape) != spatial_rank:
         raise ShapeInferenceError(
-            "AveragePool output shape must be non-negative"
+            "AveragePool kernel_shape must have "
+            f"{spatial_rank} dims, got {kernel_shape}"
         )
+    strides = tuple(
+        int(value) for value in node.attrs.get("strides", (1,) * spatial_rank)
+    )
+    if len(strides) != spatial_rank:
+        raise UnsupportedOpError("AveragePool stride rank mismatch")
+    dilations = tuple(
+        int(value)
+        for value in node.attrs.get("dilations", (1,) * spatial_rank)
+    )
+    if len(dilations) != spatial_rank:
+        raise UnsupportedOpError("AveragePool dilation rank mismatch")
+    pads = tuple(
+        int(value) for value in node.attrs.get("pads", (0,) * (2 * spatial_rank))
+    )
+    if len(pads) != 2 * spatial_rank:
+        raise UnsupportedOpError("AveragePool pads rank mismatch")
+    if auto_pad in ("", "NOTSET"):
+        pad_begin = pads[:spatial_rank]
+        pad_end = pads[spatial_rank:]
+    elif auto_pad == "VALID":
+        pad_begin = (0,) * spatial_rank
+        pad_end = (0,) * spatial_rank
+    elif auto_pad in {"SAME_UPPER", "SAME_LOWER"}:
+        pad_begin = []
+        pad_end = []
+        for dim, stride, dilation, kernel in zip(
+            input_shape[2:], strides, dilations, kernel_shape
+        ):
+            effective_kernel = dilation * (kernel - 1) + 1
+            out_dim = math.ceil(dim / stride)
+            pad_needed = max(0, (out_dim - 1) * stride + effective_kernel - dim)
+            if auto_pad == "SAME_UPPER":
+                pad_start = pad_needed // 2
+            else:
+                pad_start = (pad_needed + 1) // 2
+            pad_begin.append(pad_start)
+            pad_end.append(pad_needed - pad_start)
+        pad_begin = tuple(pad_begin)
+        pad_end = tuple(pad_end)
+    else:
+        raise UnsupportedOpError("AveragePool has unsupported auto_pad mode")
+    batch, channels = input_shape[:2]
+    in_spatial = input_shape[2:]
+    out_spatial = []
+    for dim, stride, dilation, kernel, pad_start, pad_finish in zip(
+        in_spatial, strides, dilations, kernel_shape, pad_begin, pad_end
+    ):
+        effective_kernel = dilation * (kernel - 1) + 1
+        numerator = dim + pad_start + pad_finish - effective_kernel
+        if ceil_mode:
+            out_dim = (numerator + stride - 1) // stride + 1
+            if (out_dim - 1) * stride >= dim + pad_start:
+                out_dim -= 1
+        else:
+            out_dim = numerator // stride + 1
+        if out_dim < 0:
+            raise ShapeInferenceError(
+                "AveragePool output shape must be non-negative"
+            )
+        out_spatial.append(out_dim)
     output_shape = _value_shape(graph, node.outputs[0], node)
-    expected_output_shape = (batch, channels, out_h, out_w)
+    expected_output_shape = (batch, channels, *out_spatial)
     if output_shape != expected_output_shape:
         raise ShapeInferenceError(
             "AveragePool output shape must be "
             f"{expected_output_shape}, got {output_shape}"
         )
+    in_d = in_spatial[0] if spatial_rank == 3 else 1
+    in_h = in_spatial[-2] if spatial_rank >= 2 else 1
+    in_w = in_spatial[-1]
+    out_d = out_spatial[0] if spatial_rank == 3 else 1
+    out_h = out_spatial[-2] if spatial_rank >= 2 else 1
+    out_w = out_spatial[-1]
+    kernel_d = kernel_shape[0] if spatial_rank == 3 else 1
+    kernel_h = kernel_shape[-2] if spatial_rank >= 2 else 1
+    kernel_w = kernel_shape[-1]
+    dilation_d = dilations[0] if spatial_rank == 3 else 1
+    dilation_h = dilations[-2] if spatial_rank >= 2 else 1
+    dilation_w = dilations[-1]
+    stride_d = strides[0] if spatial_rank == 3 else 1
+    stride_h = strides[-2] if spatial_rank >= 2 else 1
+    stride_w = strides[-1]
+    pad_front = pad_begin[0] if spatial_rank == 3 else 0
+    pad_top = pad_begin[-2] if spatial_rank >= 2 else 0
+    pad_left = pad_begin[-1]
+    pad_back = pad_end[0] if spatial_rank == 3 else 0
+    pad_bottom = pad_end[-2] if spatial_rank >= 2 else 0
+    pad_right = pad_end[-1]
     return _AveragePoolSpec(
         batch=batch,
         channels=channels,
+        spatial_rank=spatial_rank,
+        in_d=in_d,
         in_h=in_h,
         in_w=in_w,
+        out_d=out_d,
         out_h=out_h,
         out_w=out_w,
+        kernel_d=kernel_d,
         kernel_h=kernel_h,
         kernel_w=kernel_w,
+        dilation_d=dilation_d,
+        dilation_h=dilation_h,
+        dilation_w=dilation_w,
+        stride_d=stride_d,
         stride_h=stride_h,
         stride_w=stride_w,
+        pad_front=pad_front,
         pad_top=pad_top,
         pad_left=pad_left,
+        pad_back=pad_back,
         pad_bottom=pad_bottom,
         pad_right=pad_right,
         count_include_pad=bool(count_include_pad),
@@ -128,29 +219,48 @@ def _resolve_global_average_pool_spec(graph: Graph, node: Node) -> _AveragePoolS
     if node.attrs:
         raise UnsupportedOpError("GlobalAveragePool has unsupported attributes")
     input_shape = _value_shape(graph, node.inputs[0], node)
-    if len(input_shape) != 4:
-        raise UnsupportedOpError("GlobalAveragePool supports NCHW 2D inputs only")
-    batch, channels, in_h, in_w = input_shape
+    if len(input_shape) < 3:
+        raise UnsupportedOpError(
+            "GlobalAveragePool expects NCHW inputs with spatial dims"
+        )
+    spatial_rank = len(input_shape) - 2
+    if spatial_rank not in {1, 2, 3}:
+        raise UnsupportedOpError("GlobalAveragePool supports 1D/2D/3D inputs only")
+    batch, channels = input_shape[:2]
+    in_spatial = input_shape[2:]
     output_shape = _value_shape(graph, node.outputs[0], node)
-    expected_output_shape = (batch, channels, 1, 1)
+    expected_output_shape = (batch, channels, *([1] * spatial_rank))
     if output_shape != expected_output_shape:
         raise ShapeInferenceError(
             "GlobalAveragePool output shape must be "
             f"{expected_output_shape}, got {output_shape}"
         )
+    in_d = in_spatial[0] if spatial_rank == 3 else 1
+    in_h = in_spatial[-2] if spatial_rank >= 2 else 1
+    in_w = in_spatial[-1]
     return _AveragePoolSpec(
         batch=batch,
         channels=channels,
+        spatial_rank=spatial_rank,
+        in_d=in_d,
         in_h=in_h,
         in_w=in_w,
+        out_d=1,
         out_h=1,
         out_w=1,
+        kernel_d=in_d,
         kernel_h=in_h,
         kernel_w=in_w,
+        dilation_d=1,
+        dilation_h=1,
+        dilation_w=1,
+        stride_d=1,
         stride_h=1,
         stride_w=1,
+        pad_front=0,
         pad_top=0,
         pad_left=0,
+        pad_back=0,
         pad_bottom=0,
         pad_right=0,
         count_include_pad=False,
@@ -176,16 +286,26 @@ def lower_average_pool(graph: Graph, node: Node) -> AveragePoolOp:
         output=node.outputs[0],
         batch=spec.batch,
         channels=spec.channels,
+        spatial_rank=spec.spatial_rank,
+        in_d=spec.in_d,
         in_h=spec.in_h,
         in_w=spec.in_w,
+        out_d=spec.out_d,
         out_h=spec.out_h,
         out_w=spec.out_w,
+        kernel_d=spec.kernel_d,
         kernel_h=spec.kernel_h,
         kernel_w=spec.kernel_w,
+        dilation_d=spec.dilation_d,
+        dilation_h=spec.dilation_h,
+        dilation_w=spec.dilation_w,
+        stride_d=spec.stride_d,
         stride_h=spec.stride_h,
         stride_w=spec.stride_w,
+        pad_front=spec.pad_front,
         pad_top=spec.pad_top,
         pad_left=spec.pad_left,
+        pad_back=spec.pad_back,
         pad_bottom=spec.pad_bottom,
         pad_right=spec.pad_right,
         count_include_pad=spec.count_include_pad,
@@ -212,16 +332,26 @@ def lower_global_average_pool(graph: Graph, node: Node) -> AveragePoolOp:
         output=node.outputs[0],
         batch=spec.batch,
         channels=spec.channels,
+        spatial_rank=spec.spatial_rank,
+        in_d=spec.in_d,
         in_h=spec.in_h,
         in_w=spec.in_w,
+        out_d=spec.out_d,
         out_h=spec.out_h,
         out_w=spec.out_w,
+        kernel_d=spec.kernel_d,
         kernel_h=spec.kernel_h,
         kernel_w=spec.kernel_w,
+        dilation_d=spec.dilation_d,
+        dilation_h=spec.dilation_h,
+        dilation_w=spec.dilation_w,
+        stride_d=spec.stride_d,
         stride_h=spec.stride_h,
         stride_w=spec.stride_w,
+        pad_front=spec.pad_front,
         pad_top=spec.pad_top,
         pad_left=spec.pad_left,
+        pad_back=spec.pad_back,
         pad_bottom=spec.pad_bottom,
         pad_right=spec.pad_right,
         count_include_pad=spec.count_include_pad,

emx_onnx_cgen/lowering/bernoulli.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..dtypes import dtype_info
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..ir.ops import BernoulliOp
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+_SUPPORTED_INPUT_DTYPES = {ScalarType.F16, ScalarType.F32, ScalarType.F64}
+_SUPPORTED_OUTPUT_DTYPES = {
+    ScalarType.U8,
+    ScalarType.U16,
+    ScalarType.U32,
+    ScalarType.U64,
+    ScalarType.I8,
+    ScalarType.I16,
+    ScalarType.I32,
+    ScalarType.I64,
+    ScalarType.F16,
+    ScalarType.F32,
+    ScalarType.F64,
+    ScalarType.BOOL,
+}
+
+
+@register_lowering("Bernoulli")
+def lower_bernoulli(graph: Graph, node: Node) -> BernoulliOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Bernoulli 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 input_shape != output_shape:
+        raise ShapeInferenceError(
+            "Bernoulli output shape must match input shape, "
+            f"got {output_shape} for input {input_shape}"
+        )
+    input_dtype = _value_dtype(graph, node.inputs[0], node)
+    if input_dtype not in _SUPPORTED_INPUT_DTYPES:
+        raise UnsupportedOpError(
+            "Bernoulli input dtype must be float, "
+            f"got {input_dtype.onnx_name}"
+        )
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    dtype_attr = node.attrs.get("dtype")
+    if dtype_attr is not None:
+        attr_dtype = dtype_info(int(dtype_attr))
+        if attr_dtype != output_dtype:
+            raise UnsupportedOpError(
+                "Bernoulli dtype attribute does not match output dtype"
+            )
+    if output_dtype not in _SUPPORTED_OUTPUT_DTYPES:
+        raise UnsupportedOpError(
+            "Bernoulli output dtype must be numeric or bool, "
+            f"got {output_dtype.onnx_name}"
+        )
+    seed_value = node.attrs.get("seed")
+    seed = None
+    if seed_value is not None:
+        seed = int(seed_value)
+    return BernoulliOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        input_dtype=input_dtype,
+        dtype=output_dtype,
+        seed=seed,
+    )

emx_onnx_cgen/lowering/common.py

@@ -50,6 +50,8 @@ def value_shape(
     if isinstance(graph, GraphContext):
         shape = graph.shape(name, node)
         value = graph.find_value(name)
+        if graph.has_shape(name):
+            return shape
     else:
         try:
             value = graph.find_value(name)
@@ -219,6 +221,37 @@ def _shape_values_from_input(
                 return [int(l / r) if r != 0 else 0 for l, r in zip(left, right)]
             if source_node.op_type == "Mod":
                 return [l % r if r != 0 else 0 for l, r in zip(left, right)]
+        if source_node.op_type in {"Add", "Sub", "Mul"}:
+            if len(source_node.inputs) != 2 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError(
+                    f"{source_node.op_type} must have 2 inputs and 1 output"
+                )
+            left = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            right = _shape_values_from_input(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            if left is None or right is None:
+                return None
+            if len(left) == 1 and len(right) != 1:
+                left = left * len(right)
+            if len(right) == 1 and len(left) != 1:
+                right = right * len(left)
+            if len(left) != len(right):
+                return None
+            if source_node.op_type == "Add":
+                return [l + r for l, r in zip(left, right)]
+            if source_node.op_type == "Sub":
+                return [l - r for l, r in zip(left, right)]
+            if source_node.op_type == "Mul":
+                return [l * r for l, r in zip(left, right)]
         if source_node.op_type == "Not":
             if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
                 raise UnsupportedOpError("Not must have 1 input and 1 output")
@@ -465,3 +498,18 @@ def optional_name(names: Sequence[str], index: int) -> str | None:
         return None
     name = names[index]
     return name or None
+
+
+def resolve_int_list_from_value(
+    graph: Graph | GraphContext,
+    name: str,
+    node: Node | None = None,
+) -> list[int] | None:
+    return _shape_values_from_input(graph, name, node)
+
+
+def value_has_dim_params(
+    graph: Graph | GraphContext,
+    name: str,
+) -> bool:
+    return any(graph.find_value(name).type.dim_params)

emx_onnx_cgen/lowering/concat.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-from ..ir.ops import ConcatOp
 from ..errors import UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
+from ..ir.ops import ConcatOp
 from .common import node_dtype as _node_dtype
+from .common import value_has_dim_params as _value_has_dim_params
 from .common import value_shape as _value_shape
 from .registry import register_lowering
-from ..validation import validate_concat_shapes
+from ..validation import normalize_concat_axis, validate_concat_shapes
 
 
 @register_lowering("Concat")
@@ -15,12 +17,44 @@ def lower_concat(graph: Graph, node: Node) -> ConcatOp:
         raise UnsupportedOpError("Concat must have at least 1 input and 1 output")
     op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
     output_shape = _value_shape(graph, node.outputs[0], node)
+    if _value_has_dim_params(graph, node.outputs[0]):
+        output_shape = ()
     input_shapes = tuple(_value_shape(graph, name, node) for name in node.inputs)
-    axis = validate_concat_shapes(
-        input_shapes,
-        output_shape,
-        int(node.attrs.get("axis", 0)),
-    )
+    axis = int(node.attrs.get("axis", 0))
+    if output_shape:
+        axis = validate_concat_shapes(
+            input_shapes,
+            output_shape,
+            axis,
+        )
+    else:
+        ranks = {len(shape) for shape in input_shapes}
+        if len(ranks) != 1:
+            raise UnsupportedOpError(
+                f"Concat inputs must have matching ranks, got {input_shapes}"
+            )
+        rank = ranks.pop()
+        axis = normalize_concat_axis(axis, rank)
+        base_shape = list(input_shapes[0])
+        axis_dim = 0
+        for shape in input_shapes:
+            if len(shape) != rank:
+                raise UnsupportedOpError(
+                    f"Concat inputs must have matching ranks, got {input_shapes}"
+                )
+            for dim_index, dim in enumerate(shape):
+                if dim_index == axis:
+                    continue
+                if dim != base_shape[dim_index]:
+                    raise UnsupportedOpError(
+                        "Concat inputs must match on non-axis dimensions, "
+                        f"got {input_shapes}"
+                    )
+            axis_dim += shape[axis]
+        base_shape[axis] = axis_dim
+        output_shape = tuple(base_shape)
+        if isinstance(graph, GraphContext):
+            graph.set_shape(node.outputs[0], output_shape)
     return ConcatOp(
         inputs=node.inputs,
         output=node.outputs[0],

emx_onnx_cgen/lowering/conv.py

@@ -26,9 +26,14 @@ class ConvSpec:
     group: int
 
 
-def resolve_conv_spec(graph: Graph, node: Node) -> ConvSpec:
-    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
-        raise UnsupportedOpError("Conv must have 2 or 3 inputs and 1 output")
+def resolve_conv_spec(
+    graph: Graph,
+    node: Node,
+    *,
+    input_name: str,
+    weight_name: str,
+    bias_name: str | None,
+) -> ConvSpec:
     supported_attrs = {
         "auto_pad",
         "dilations",
@@ -39,8 +44,8 @@ def resolve_conv_spec(graph: Graph, node: Node) -> ConvSpec:
     }
     if set(node.attrs) - supported_attrs:
         raise UnsupportedOpError("Conv has unsupported attributes")
-    input_shape = _value_shape(graph, node.inputs[0], node)
-    weight_shape = _value_shape(graph, node.inputs[1], node)
+    input_shape = _value_shape(graph, input_name, node)
+    weight_shape = _value_shape(graph, weight_name, node)
     if len(input_shape) < 3:
         raise UnsupportedOpError("Conv expects NCHW inputs with spatial dims")
     spatial_rank = len(input_shape) - 2
@@ -79,8 +84,8 @@ def resolve_conv_spec(graph: Graph, node: Node) -> ConvSpec:
             "Conv input channels must match weight channels, "
             f"got {in_channels} and {weight_in_channels * group}"
         )
-    if len(node.inputs) == 3:
-        bias_shape = _value_shape(graph, node.inputs[2], node)
+    if bias_name is not None:
+        bias_shape = _value_shape(graph, bias_name, node)
         if bias_shape != (out_channels,):
             raise ShapeInferenceError(
                 f"Conv bias shape must be {(out_channels,)}, got {bias_shape}"
@@ -171,7 +176,13 @@ def lower_conv(graph: Graph, node: Node) -> ConvOp:
         raise UnsupportedOpError(
             "Conv supports float16, float, and double inputs only"
         )
-    spec = resolve_conv_spec(graph, node)
+    spec = resolve_conv_spec(
+        graph,
+        node,
+        input_name=node.inputs[0],
+        weight_name=node.inputs[1],
+        bias_name=node.inputs[2] if len(node.inputs) == 3 else None,
+    )
     return ConvOp(
         input0=node.inputs[0],
         weights=node.inputs[1],

emx_onnx_cgen/lowering/conv_integer.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..errors import UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..ir.ops import ConvIntegerOp
+from .common import optional_name, value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .conv import resolve_conv_spec
+from .registry import register_lowering
+
+
+def _ensure_scalar_shape(shape: tuple[int, ...], label: str) -> None:
+    if shape not in {(), (1,)}:
+        raise UnsupportedOpError(
+            f"ConvInteger {label} must be a scalar, got shape {shape}"
+        )
+
+
+def _resolve_w_zero_point_shape(
+    shape: tuple[int, ...], out_channels: int
+) -> bool:
+    if shape in {(), (1,)}:
+        return False
+    if shape == (out_channels,):
+        return True
+    raise UnsupportedOpError(
+        "ConvInteger w_zero_point must be scalar or 1D per output channel, "
+        f"got shape {shape}"
+    )
+
+
+@register_lowering("ConvInteger")
+def lower_conv_integer(graph: Graph, node: Node) -> ConvIntegerOp:
+    if len(node.inputs) not in {2, 3, 4} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "ConvInteger must have 2 to 4 inputs and 1 output"
+        )
+    input_name = node.inputs[0]
+    weight_name = node.inputs[1]
+    x_zero_point_name = optional_name(node.inputs, 2)
+    w_zero_point_name = optional_name(node.inputs, 3)
+    input_dtype = _value_dtype(graph, input_name, node)
+    weight_dtype = _value_dtype(graph, weight_name, node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if input_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError("ConvInteger supports uint8/int8 inputs only")
+    if weight_dtype not in {ScalarType.U8, ScalarType.I8}:
+        raise UnsupportedOpError("ConvInteger supports uint8/int8 weights only")
+    if output_dtype != ScalarType.I32:
+        raise UnsupportedOpError("ConvInteger expects int32 outputs only")
+    x_zero_shape = None
+    if x_zero_point_name is not None:
+        x_zero_shape = _value_shape(graph, x_zero_point_name, node)
+        _ensure_scalar_shape(x_zero_shape, "x_zero_point")
+        if _value_dtype(graph, x_zero_point_name, node) != input_dtype:
+            raise UnsupportedOpError(
+                "ConvInteger x_zero_point dtype must match input dtype"
+            )
+    w_zero_shape = None
+    w_zero_point_per_channel = False
+    if w_zero_point_name is not None:
+        w_zero_shape = _value_shape(graph, w_zero_point_name, node)
+        if _value_dtype(graph, w_zero_point_name, node) != weight_dtype:
+            raise UnsupportedOpError(
+                "ConvInteger w_zero_point dtype must match weight dtype"
+            )
+    spec = resolve_conv_spec(
+        graph,
+        node,
+        input_name=input_name,
+        weight_name=weight_name,
+        bias_name=None,
+    )
+    if w_zero_shape is not None:
+        w_zero_point_per_channel = _resolve_w_zero_point_shape(
+            w_zero_shape, spec.out_channels
+        )
+    return ConvIntegerOp(
+        input0=input_name,
+        weights=weight_name,
+        x_zero_point=x_zero_point_name,
+        w_zero_point=w_zero_point_name,
+        output=node.outputs[0],
+        batch=spec.batch,
+        in_channels=spec.in_channels,
+        out_channels=spec.out_channels,
+        spatial_rank=spec.spatial_rank,
+        in_spatial=spec.in_spatial,
+        out_spatial=spec.out_spatial,
+        kernel_shape=spec.kernel_shape,
+        strides=spec.strides,
+        pads=spec.pads,
+        dilations=spec.dilations,
+        group=spec.group,
+        input_dtype=input_dtype,
+        weight_dtype=weight_dtype,
+        dtype=output_dtype,
+        x_zero_point_shape=x_zero_shape,
+        w_zero_point_shape=w_zero_shape,
+        w_zero_point_per_channel=w_zero_point_per_channel,
+    )