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

emx_onnx_cgen/lowering/einsum.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import EinsumKind, EinsumOp
+from ..errors import ShapeInferenceError, 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
+
+
+def _normalize_equation(equation: str) -> str:
+    return equation.replace(" ", "")
+
+
+@register_lowering("Einsum")
+def lower_einsum(graph: Graph, node: Node) -> EinsumOp:
+    if not node.inputs or len(node.outputs) != 1:
+        raise UnsupportedOpError("Einsum must have 1 output and at least 1 input")
+    equation_value = node.attrs.get("equation")
+    if equation_value is None:
+        raise UnsupportedOpError("Einsum equation attribute is required")
+    equation = (
+        equation_value.decode()
+        if isinstance(equation_value, (bytes, bytearray))
+        else str(equation_value)
+    )
+    normalized = _normalize_equation(equation)
+    input_shapes = tuple(
+        _value_shape(graph, name, node) for name in node.inputs
+    )
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if normalized == "->":
+        if len(node.inputs) != 1:
+            raise UnsupportedOpError("Einsum '->' must have 1 input")
+        if output_shape:
+            raise ShapeInferenceError(
+                "Einsum '->' output must be scalar, "
+                f"got shape {output_shape}"
+            )
+        kind = EinsumKind.REDUCE_ALL
+    elif normalized == "ij->i":
+        if len(node.inputs) != 1:
+            raise UnsupportedOpError("Einsum 'ij->i' must have 1 input")
+        input_shape = input_shapes[0]
+        if len(input_shape) != 2:
+            raise ShapeInferenceError(
+                "Einsum 'ij->i' input must be 2D, "
+                f"got shape {input_shape}"
+            )
+        expected = (input_shape[0],)
+        if output_shape != expected:
+            raise ShapeInferenceError(
+                f"Einsum 'ij->i' output must match shape {expected}, "
+                f"got {output_shape}"
+            )
+        kind = EinsumKind.SUM_J
+    elif normalized == "ij->ji":
+        if len(node.inputs) != 1:
+            raise UnsupportedOpError("Einsum 'ij->ji' must have 1 input")
+        input_shape = input_shapes[0]
+        if len(input_shape) != 2:
+            raise ShapeInferenceError(
+                "Einsum 'ij->ji' input must be 2D, "
+                f"got shape {input_shape}"
+            )
+        expected = (input_shape[1], input_shape[0])
+        if output_shape != expected:
+            raise ShapeInferenceError(
+                f"Einsum 'ij->ji' output must match shape {expected}, "
+                f"got {output_shape}"
+            )
+        kind = EinsumKind.TRANSPOSE
+    elif normalized in {"i,i", "i,i->"}:
+        if len(node.inputs) != 2:
+            raise UnsupportedOpError("Einsum 'i,i' must have 2 inputs")
+        left_shape, right_shape = input_shapes
+        if len(left_shape) != 1 or len(right_shape) != 1:
+            raise ShapeInferenceError(
+                "Einsum 'i,i' inputs must be vectors, "
+                f"got shapes {left_shape} and {right_shape}"
+            )
+        if left_shape[0] != right_shape[0]:
+            raise ShapeInferenceError(
+                "Einsum 'i,i' inputs must have the same length, "
+                f"got shapes {left_shape} and {right_shape}"
+            )
+        if output_shape:
+            raise ShapeInferenceError(
+                "Einsum 'i,i' output must be scalar, "
+                f"got shape {output_shape}"
+            )
+        kind = EinsumKind.DOT
+    elif normalized == "bij,bjk->bik":
+        if len(node.inputs) != 2:
+            raise UnsupportedOpError("Einsum 'bij,bjk->bik' must have 2 inputs")
+        left_shape, right_shape = input_shapes
+        if len(left_shape) != 3 or len(right_shape) != 3:
+            raise ShapeInferenceError(
+                "Einsum 'bij,bjk->bik' inputs must be 3D, "
+                f"got shapes {left_shape} and {right_shape}"
+            )
+        if left_shape[0] != right_shape[0]:
+            raise ShapeInferenceError(
+                "Einsum 'bij,bjk->bik' batch dimensions must match, "
+                f"got shapes {left_shape} and {right_shape}"
+            )
+        if left_shape[2] != right_shape[1]:
+            raise ShapeInferenceError(
+                "Einsum 'bij,bjk->bik' contraction dimensions must match, "
+                f"got shapes {left_shape} and {right_shape}"
+            )
+        expected = (left_shape[0], left_shape[1], right_shape[2])
+        if output_shape != expected:
+            raise ShapeInferenceError(
+                f"Einsum 'bij,bjk->bik' output must match shape {expected}, "
+                f"got {output_shape}"
+            )
+        kind = EinsumKind.BATCH_MATMUL
+    elif normalized == "...ii->...i":
+        if len(node.inputs) != 1:
+            raise UnsupportedOpError("Einsum '...ii->...i' must have 1 input")
+        input_shape = input_shapes[0]
+        if len(input_shape) < 2:
+            raise ShapeInferenceError(
+                "Einsum '...ii->...i' input must be at least 2D, "
+                f"got shape {input_shape}"
+            )
+        if input_shape[-1] != input_shape[-2]:
+            raise ShapeInferenceError(
+                "Einsum '...ii->...i' requires last two dims to match, "
+                f"got shape {input_shape}"
+            )
+        expected = (*input_shape[:-2], input_shape[-1])
+        if output_shape != expected:
+            raise ShapeInferenceError(
+                f"Einsum '...ii->...i' output must match shape {expected}, "
+                f"got {output_shape}"
+            )
+        kind = EinsumKind.BATCH_DIAGONAL
+    else:
+        raise UnsupportedOpError(
+            f"Unsupported Einsum equation '{equation}'"
+        )
+    return EinsumOp(
+        inputs=tuple(node.inputs),
+        output=node.outputs[0],
+        kind=kind,
+        input_shapes=input_shapes,
+        output_shape=output_shape,
+        dtype=op_dtype,
+        input_dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/gather_elements.py

@@ -33,9 +33,7 @@ def lower_gather_elements(graph: Graph, node: Node) -> GatherElementsOp:
     for dim_index, (data_dim, index_dim) in enumerate(
         zip(data_shape, indices_shape)
     ):
-        if dim_index == axis:
-            continue
-        if data_dim != index_dim:
+        if dim_index != axis and data_dim != index_dim:
             raise ShapeInferenceError(
                 "GatherElements inputs must match on non-axis dimensions, "
                 f"got {data_shape} and {indices_shape}"

emx_onnx_cgen/lowering/gather_nd.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import GatherNDOp
+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
+
+
+@register_lowering("GatherND")
+def lower_gather_nd(graph: Graph, node: Node) -> GatherNDOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("GatherND must have 2 inputs and 1 output")
+    data_name, indices_name = node.inputs
+    output_name = node.outputs[0]
+    data_shape = _value_shape(graph, data_name, node)
+    indices_shape = _value_shape(graph, indices_name, node)
+    output_shape = _value_shape(graph, output_name, node)
+    if len(indices_shape) < 1:
+        raise ShapeInferenceError("GatherND indices must have rank >= 1")
+    batch_dims = int(node.attrs.get("batch_dims", 0))
+    if batch_dims < 0:
+        raise ShapeInferenceError(
+            f"GatherND batch_dims must be >= 0, got {batch_dims}"
+        )
+    if batch_dims > len(indices_shape) - 1:
+        raise ShapeInferenceError(
+            "GatherND batch_dims must be <= indices rank - 1, "
+            f"got {batch_dims} vs {len(indices_shape) - 1}"
+        )
+    if batch_dims > len(data_shape):
+        raise ShapeInferenceError(
+            "GatherND batch_dims must be <= data rank, "
+            f"got {batch_dims} vs {len(data_shape)}"
+        )
+    if tuple(data_shape[:batch_dims]) != tuple(indices_shape[:batch_dims]):
+        raise ShapeInferenceError(
+            "GatherND batch_dims must match on data/indices, "
+            f"got {data_shape} vs {indices_shape}"
+        )
+    index_depth = indices_shape[-1]
+    if index_depth <= 0:
+        raise ShapeInferenceError(
+            "GatherND indices final dimension must be >= 1"
+        )
+    if index_depth > len(data_shape) - batch_dims:
+        raise ShapeInferenceError(
+            "GatherND indices final dimension must be <= data rank - "
+            f"batch_dims, got {index_depth} vs {len(data_shape) - batch_dims}"
+        )
+    expected_output_shape = indices_shape[:-1] + data_shape[
+        batch_dims + index_depth :
+    ]
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "GatherND output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    data_dtype = _value_dtype(graph, data_name, node)
+    indices_dtype = _value_dtype(graph, indices_name, node)
+    if indices_dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            "GatherND indices must be int32 or int64, "
+            f"got {indices_dtype.onnx_name}"
+        )
+    return GatherNDOp(
+        data=data_name,
+        indices=indices_name,
+        output=output_name,
+        batch_dims=batch_dims,
+        data_shape=data_shape,
+        indices_shape=indices_shape,
+        output_shape=output_shape,
+        dtype=data_dtype,
+        indices_dtype=indices_dtype,
+    )

emx_onnx_cgen/lowering/global_max_pool.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import ReduceOp
+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
+
+
+@register_lowering("GlobalMaxPool")
+def lower_global_max_pool(graph: Graph, node: Node) -> ReduceOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("GlobalMaxPool must have 1 input and 1 output")
+    if node.attrs:
+        raise UnsupportedOpError("GlobalMaxPool has unsupported attributes")
+    op_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if op_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "GlobalMaxPool expects matching input/output dtypes, "
+            f"got {op_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if op_dtype not in {ScalarType.F16, ScalarType.F32, ScalarType.F64}:
+        raise UnsupportedOpError(
+            "GlobalMaxPool supports float16, float, and double inputs only"
+        )
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) < 3:
+        raise UnsupportedOpError(
+            "GlobalMaxPool expects input rank of at least 3"
+        )
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    expected_output_shape = (input_shape[0], input_shape[1]) + (
+        1,
+    ) * (len(input_shape) - 2)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "GlobalMaxPool output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    axes = tuple(range(2, len(input_shape)))
+    return ReduceOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        axes=axes,
+        axes_input=None,
+        axes_input_shape=None,
+        axes_input_dtype=None,
+        keepdims=True,
+        noop_with_empty_axes=False,
+        reduce_kind="max",
+        reduce_count=None,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/hardmax.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import HardmaxOp
+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("Hardmax")
+def lower_hardmax(graph: Graph, node: Node) -> HardmaxOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Hardmax must have 1 input and 1 output")
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if op_dtype not in {ScalarType.F16, ScalarType.F32, ScalarType.F64}:
+        raise UnsupportedOpError(
+            "Hardmax 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 HardmaxOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        outer=outer,
+        axis_size=axis_size,
+        inner=inner,
+        axis=axis,
+        shape=input_shape,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/identity.py

@@ -22,11 +22,12 @@ def lower_identity(graph: Graph, node: Node) -> IdentityOp:
         for index, (input_dim, output_dim) in enumerate(
             zip(input_shape, output_shape)
         ):
-            if input_dim == output_dim:
-                continue
-            if input_dim_params[index] or output_dim_params[index]:
-                continue
-            raise ShapeInferenceError("Identity input and output shapes must match")
+            if input_dim != output_dim and not (
+                input_dim_params[index] or output_dim_params[index]
+            ):
+                raise ShapeInferenceError(
+                    "Identity input and output shapes must match"
+                )
     input_dtype = value_dtype(graph, node.inputs[0], node)
     output_dtype = value_dtype(graph, node.outputs[0], node)
     if input_dtype != output_dtype:

emx_onnx_cgen/lowering/logsoftmax.py

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

emx_onnx_cgen/lowering/lp_pool.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import LpPoolOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .registry import register_lowering
+from .common import value_dtype as _value_dtype, value_shape as _value_shape
+
+
+@dataclass(frozen=True)
+class LpPoolSpec:
+    batch: int
+    channels: int
+    in_h: int
+    in_w: int
+    out_h: int
+    out_w: int
+    kernel_h: int
+    kernel_w: int
+    stride_h: int
+    stride_w: int
+    pad_top: int
+    pad_left: int
+    pad_bottom: int
+    pad_right: int
+    p: int
+
+
+def _resolve_lp_pool_spec(graph: Graph, node: Node) -> LpPoolSpec:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("LpPool must have 1 input and 1 output")
+    supported_attrs = {
+        "auto_pad",
+        "ceil_mode",
+        "dilations",
+        "kernel_shape",
+        "pads",
+        "p",
+        "strides",
+    }
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("LpPool has unsupported attributes")
+    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("LpPool supports auto_pad=NOTSET only")
+    ceil_mode = int(node.attrs.get("ceil_mode", 0))
+    if ceil_mode != 0:
+        raise UnsupportedOpError("LpPool supports ceil_mode=0 only")
+    dilations = tuple(int(value) for value in node.attrs.get("dilations", (1, 1)))
+    if any(value != 1 for value in dilations):
+        raise UnsupportedOpError("LpPool supports dilations=1 only")
+    kernel_shape = node.attrs.get("kernel_shape")
+    if kernel_shape is None:
+        raise UnsupportedOpError("LpPool requires kernel_shape")
+    kernel_shape = tuple(int(value) for value in kernel_shape)
+    if len(kernel_shape) != 2:
+        raise UnsupportedOpError("LpPool 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("LpPool expects 2D strides")
+    pads = tuple(int(value) for value in node.attrs.get("pads", (0, 0, 0, 0)))
+    if len(pads) != 4:
+        raise UnsupportedOpError("LpPool expects 4D pads")
+    pad_top, pad_left, pad_bottom, pad_right = pads
+    p = int(node.attrs.get("p", 2))
+    if p < 1:
+        raise UnsupportedOpError("LpPool p must be >= 1")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) != 4:
+        raise UnsupportedOpError("LpPool 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:
+        raise ShapeInferenceError("LpPool output shape must be non-negative")
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    expected_output_shape = (batch, channels, out_h, out_w)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "LpPool output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    return LpPoolSpec(
+        batch=batch,
+        channels=channels,
+        in_h=in_h,
+        in_w=in_w,
+        out_h=out_h,
+        out_w=out_w,
+        kernel_h=kernel_h,
+        kernel_w=kernel_w,
+        stride_h=stride_h,
+        stride_w=stride_w,
+        pad_top=pad_top,
+        pad_left=pad_left,
+        pad_bottom=pad_bottom,
+        pad_right=pad_right,
+        p=p,
+    )
+
+
+@register_lowering("LpPool")
+def lower_lp_pool(graph: Graph, node: Node) -> LpPoolOp:
+    op_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if op_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "LpPool expects matching input/output dtypes, "
+            f"got {op_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "LpPool supports float16, float, and double inputs only"
+        )
+    spec = _resolve_lp_pool_spec(graph, node)
+    return LpPoolOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        batch=spec.batch,
+        channels=spec.channels,
+        in_h=spec.in_h,
+        in_w=spec.in_w,
+        out_h=spec.out_h,
+        out_w=spec.out_w,
+        kernel_h=spec.kernel_h,
+        kernel_w=spec.kernel_w,
+        stride_h=spec.stride_h,
+        stride_w=spec.stride_w,
+        pad_top=spec.pad_top,
+        pad_left=spec.pad_left,
+        pad_bottom=spec.pad_bottom,
+        pad_right=spec.pad_right,
+        p=spec.p,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/matmul.py

@@ -87,13 +87,12 @@ def _broadcast_batch_shapes(
     right_padded = (1,) * (max_rank - len(right)) + right
     broadcast_shape = []
     for left_dim, right_dim in zip(left_padded, right_padded):
-        if left_dim == right_dim or left_dim == 1 or right_dim == 1:
-            broadcast_shape.append(max(left_dim, right_dim))
-            continue
-        raise ShapeInferenceError(
-            "MatMul batch dimensions must be broadcastable, "
-            f"got {left} x {right}"
-        )
+        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
 
 

emx_onnx_cgen/lowering/negative_log_likelihood_loss.py

@@ -43,18 +43,18 @@ def _resolve_target_shape(
                 raise ShapeInferenceError("Reshape allows only one -1 dimension")
             unknown_index = index
             output_dims.append(-1)
-            continue
-        if dim == 0:
-            if allowzero == 0:
-                if index >= len(input_shape):
-                    raise ShapeInferenceError(
-                        "Reshape zero dim must index into input shape"
-                    )
-                dim = input_shape[index]
-        if dim < 0:
-            raise ShapeInferenceError("Reshape dims must be >= -1")
-        output_dims.append(dim)
-        known_product *= dim
+        else:
+            if dim == 0:
+                if allowzero == 0:
+                    if index >= len(input_shape):
+                        raise ShapeInferenceError(
+                            "Reshape zero dim must index into input shape"
+                        )
+                    dim = input_shape[index]
+            if dim < 0:
+                raise ShapeInferenceError("Reshape dims must be >= -1")
+            output_dims.append(dim)
+            known_product *= dim
     input_product = _shape_product(input_shape)
     if unknown_index is not None:
         if known_product == 0 or input_product % known_product != 0:

emx_onnx_cgen/lowering/nonzero.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter 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,
+    )