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

emx_onnx_cgen/lowering/conv.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import ConvOp
+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
+
+
+@dataclass(frozen=True)
+class ConvSpec:
+    batch: int
+    in_channels: int
+    out_channels: int
+    spatial_rank: int
+    in_spatial: tuple[int, ...]
+    out_spatial: tuple[int, ...]
+    kernel_shape: tuple[int, ...]
+    strides: tuple[int, ...]
+    pads: tuple[int, ...]
+    dilations: tuple[int, ...]
+    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")
+    supported_attrs = {
+        "auto_pad",
+        "dilations",
+        "group",
+        "kernel_shape",
+        "pads",
+        "strides",
+    }
+    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)
+    if len(input_shape) < 3:
+        raise UnsupportedOpError("Conv expects NCHW inputs with spatial dims")
+    spatial_rank = len(input_shape) - 2
+    if spatial_rank not in {1, 2, 3}:
+        raise UnsupportedOpError("Conv supports 1D/2D/3D inputs only")
+    if len(weight_shape) != spatial_rank + 2:
+        raise UnsupportedOpError("Conv weight rank must match spatial rank")
+    batch, in_channels = input_shape[0], input_shape[1]
+    in_spatial = input_shape[2:]
+    out_channels, weight_in_channels, *kernel_shape = weight_shape
+    kernel_shape = node.attrs.get("kernel_shape")
+    if kernel_shape is not None:
+        kernel_shape = tuple(int(value) for value in kernel_shape)
+        if len(kernel_shape) != spatial_rank:
+            raise UnsupportedOpError(
+                "Conv kernel_shape rank must match input spatial rank"
+            )
+        if kernel_shape != tuple(weight_shape[2:]):
+            raise ShapeInferenceError(
+                "Conv kernel_shape must match weights, "
+                f"got {kernel_shape} and {tuple(weight_shape[2:])}"
+            )
+    else:
+        kernel_shape = tuple(weight_shape[2:])
+    group = int(node.attrs.get("group", 1))
+    if group <= 0:
+        raise UnsupportedOpError("Conv expects group >= 1")
+    if in_channels % group != 0 or out_channels % group != 0:
+        raise ShapeInferenceError(
+            "Conv expects group to evenly divide in/out channels, "
+            f"got group={group}, in_channels={in_channels}, "
+            f"out_channels={out_channels}"
+        )
+    if weight_in_channels != in_channels // group:
+        raise ShapeInferenceError(
+            "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_shape != (out_channels,):
+            raise ShapeInferenceError(
+                f"Conv bias shape must be {(out_channels,)}, got {bias_shape}"
+            )
+    strides = tuple(
+        int(value) for value in node.attrs.get("strides", (1,) * spatial_rank)
+    )
+    if len(strides) != spatial_rank:
+        raise UnsupportedOpError("Conv stride rank mismatch")
+    dilations = tuple(
+        int(value) for value in node.attrs.get("dilations", (1,) * spatial_rank)
+    )
+    if len(dilations) != spatial_rank:
+        raise UnsupportedOpError("Conv 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("Conv pads rank mismatch")
+    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 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(
+            in_spatial, 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("Conv has unsupported auto_pad mode")
+    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
+        out_dim = (dim + pad_start + pad_finish - effective_kernel) // stride + 1
+        if out_dim < 0:
+            raise ShapeInferenceError("Conv output shape must be non-negative")
+        out_spatial.append(out_dim)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    expected_output_shape = (batch, out_channels, *out_spatial)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "Conv output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    return ConvSpec(
+        batch=batch,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        spatial_rank=spatial_rank,
+        in_spatial=in_spatial,
+        out_spatial=tuple(out_spatial),
+        kernel_shape=kernel_shape,
+        strides=strides,
+        pads=(*pad_begin, *pad_end),
+        dilations=dilations,
+        group=group,
+    )
+
+
+@register_lowering("Conv")
+def lower_conv(graph: Graph, node: Node) -> ConvOp:
+    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")
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "Conv supports float16, float, and double inputs only"
+        )
+    spec = resolve_conv_spec(graph, node)
+    return ConvOp(
+        input0=node.inputs[0],
+        weights=node.inputs[1],
+        bias=node.inputs[2] if len(node.inputs) == 3 else None,
+        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,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/cumsum.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import CumSumOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import value_dtype, value_shape
+from ..validation import ensure_output_shape_matches_input, normalize_axis
+from .registry import register_lowering
+
+
+_SUPPORTED_CUMSUM_DTYPES = {
+    ScalarType.F16,
+    ScalarType.F32,
+    ScalarType.F64,
+    ScalarType.I32,
+    ScalarType.I64,
+    ScalarType.U32,
+    ScalarType.U64,
+}
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+def _validate_static_shape(shape: tuple[int, ...], node: Node) -> None:
+    for dim in shape:
+        if dim < 0:
+            raise ShapeInferenceError(
+                f"{node.op_type} does not support dynamic dims"
+            )
+
+
+def _read_axis_initializer(
+    initializer: Initializer, node: Node
+) -> int:
+    if initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            f"{node.op_type} axis input must be int64 or int32"
+        )
+    axis_data = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if axis_data.size != 1:
+        raise UnsupportedOpError(f"{node.op_type} axis input must be scalar")
+    return int(axis_data[0])
+
+
+@register_lowering("CumSum")
+def lower_cumsum(graph: Graph, node: Node) -> CumSumOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("CumSum must have 2 inputs and 1 output")
+    input_name = node.inputs[0]
+    axis_name = node.inputs[1]
+    if not input_name or not axis_name:
+        raise UnsupportedOpError("CumSum requires input and axis values")
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    _validate_static_shape(input_shape, node)
+    ensure_output_shape_matches_input(node, input_shape, output_shape)
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "CumSum expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if input_dtype not in _SUPPORTED_CUMSUM_DTYPES:
+        raise UnsupportedOpError(
+            f"CumSum does not support dtype {input_dtype.onnx_name}"
+        )
+    axis_initializer = _find_initializer(graph, axis_name)
+    axis_value = None
+    axis_input = None
+    axis_input_dtype = None
+    if axis_initializer is not None:
+        axis_value = normalize_axis(
+            _read_axis_initializer(axis_initializer, node),
+            input_shape,
+            node,
+        )
+    else:
+        axis_shape = value_shape(graph, axis_name, node)
+        if not _is_scalar_shape(axis_shape):
+            raise UnsupportedOpError("CumSum axis input must be scalar")
+        axis_input_dtype = value_dtype(graph, axis_name, node)
+        if axis_input_dtype not in {ScalarType.I64, ScalarType.I32}:
+            raise UnsupportedOpError(
+                "CumSum axis input must be int64 or int32"
+            )
+        axis_input = axis_name
+    exclusive = int(node.attrs.get("exclusive", 0))
+    reverse = int(node.attrs.get("reverse", 0))
+    if exclusive not in {0, 1}:
+        raise UnsupportedOpError("CumSum exclusive must be 0 or 1")
+    if reverse not in {0, 1}:
+        raise UnsupportedOpError("CumSum reverse must be 0 or 1")
+    return CumSumOp(
+        input0=input_name,
+        axis_input=axis_input,
+        axis_input_dtype=axis_input_dtype,
+        axis=axis_value,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        dtype=input_dtype,
+        input_dtype=input_dtype,
+        exclusive=bool(exclusive),
+        reverse=bool(reverse),
+    )

emx_onnx_cgen/lowering/depth_space.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import DepthToSpaceOp, SpaceToDepthOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..lowering.common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _blocksize(node: Node) -> int:
+    if "blocksize" not in node.attrs:
+        raise UnsupportedOpError(f"{node.op_type} requires blocksize attribute")
+    blocksize = int(node.attrs["blocksize"])
+    if blocksize <= 0:
+        raise UnsupportedOpError(
+            f"{node.op_type} blocksize must be > 0, got {blocksize}"
+        )
+    return blocksize
+
+
+@register_lowering("DepthToSpace")
+def lower_depth_to_space(graph: Graph, node: Node) -> DepthToSpaceOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("DepthToSpace must have 1 input and 1 output")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if len(input_shape) != 4 or len(output_shape) != 4:
+        raise UnsupportedOpError("DepthToSpace only supports 4D inputs")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "DepthToSpace expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    blocksize = _blocksize(node)
+    mode_attr = node.attrs.get("mode", "DCR")
+    if isinstance(mode_attr, bytes):
+        mode = mode_attr.decode()
+    else:
+        mode = str(mode_attr)
+    if mode not in {"DCR", "CRD"}:
+        raise UnsupportedOpError(
+            "DepthToSpace only supports mode DCR or CRD"
+        )
+    n, c, h, w = input_shape
+    if c % (blocksize * blocksize) != 0:
+        raise ShapeInferenceError(
+            "DepthToSpace input channels must be divisible by blocksize^2"
+        )
+    expected_shape = (
+        n,
+        c // (blocksize * blocksize),
+        h * blocksize,
+        w * blocksize,
+    )
+    if output_shape != expected_shape:
+        raise ShapeInferenceError(
+            "DepthToSpace output shape mismatch: "
+            f"expected {expected_shape}, got {output_shape}"
+        )
+    return DepthToSpaceOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        blocksize=blocksize,
+        mode=mode,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )
+
+
+@register_lowering("SpaceToDepth")
+def lower_space_to_depth(graph: Graph, node: Node) -> SpaceToDepthOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("SpaceToDepth must have 1 input and 1 output")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if len(input_shape) != 4 or len(output_shape) != 4:
+        raise UnsupportedOpError("SpaceToDepth only supports 4D inputs")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "SpaceToDepth expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    blocksize = _blocksize(node)
+    n, c, h, w = input_shape
+    if h % blocksize != 0 or w % blocksize != 0:
+        raise ShapeInferenceError(
+            "SpaceToDepth spatial dims must be divisible by blocksize"
+        )
+    expected_shape = (
+        n,
+        c * blocksize * blocksize,
+        h // blocksize,
+        w // blocksize,
+    )
+    if output_shape != expected_shape:
+        raise ShapeInferenceError(
+            "SpaceToDepth output shape mismatch: "
+            f"expected {expected_shape}, got {output_shape}"
+        )
+    return SpaceToDepthOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        blocksize=blocksize,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/dropout.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import ReshapeOp
+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
+
+
+def _is_value_used(graph: Graph, name: str) -> bool:
+    if any(value.name == name for value in graph.outputs):
+        return True
+    return any(name in node.inputs for node in graph.nodes)
+
+
+@register_lowering("Dropout")
+def lower_dropout(graph: Graph, node: Node) -> ReshapeOp:
+    if len(node.outputs) not in {1, 2} or len(node.inputs) != 1:
+        raise UnsupportedOpError(
+            "Dropout supports only the data input and 1 or 2 outputs"
+        )
+    if len(node.outputs) == 2 and _is_value_used(graph, node.outputs[1]):
+        raise UnsupportedOpError("Dropout mask output is not supported")
+    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(
+            "Dropout output shape must match input shape, "
+            f"got {output_shape} for input {input_shape}"
+        )
+    input_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Dropout expects matching input/output dtypes, "
+            f"got {input_dtype} and {output_dtype}"
+        )
+    return ReshapeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        dtype=input_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/elementwise.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+from shared.scalar_functions import ScalarFunction
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import ClipOp, UnaryOp
+from ..errors import UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..lowering.common import node_dtype, optional_name, value_dtype, value_shape
+from ..lowering.registry import register_lowering
+
+
+@register_lowering("Clip")
+def lower_clip(graph: Graph, node: Node) -> ClipOp:
+    if not node.inputs or len(node.outputs) != 1:
+        raise UnsupportedOpError("Clip must have 1 output")
+    input_name = node.inputs[0]
+    if not input_name:
+        raise UnsupportedOpError("Clip input must be provided")
+    min_name = optional_name(node.inputs, 1)
+    max_name = optional_name(node.inputs, 2)
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Clip expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if min_name is not None:
+        min_dtype = value_dtype(graph, min_name, node)
+        if min_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "Clip min dtype must match input dtype, "
+                f"got {min_dtype.onnx_name}"
+            )
+    if max_name is not None:
+        max_dtype = value_dtype(graph, max_name, node)
+        if max_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "Clip max dtype must match input dtype, "
+                f"got {max_dtype.onnx_name}"
+            )
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if input_shape != output_shape:
+        raise UnsupportedOpError("Clip input and output shapes must match")
+    min_shape = value_shape(graph, min_name, node) if min_name else None
+    max_shape = value_shape(graph, max_name, node) if max_name else None
+    return ClipOp(
+        input0=input_name,
+        input_min=min_name,
+        input_max=max_name,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        min_shape=min_shape,
+        max_shape=max_shape,
+        output_shape=output_shape,
+        dtype=input_dtype,
+    )
+
+
+@register_lowering("Celu")
+def lower_celu(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Celu must have 1 input and 1 output")
+    dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not dtype.is_float:
+        raise UnsupportedOpError("Celu only supports floating-point inputs")
+    alpha = float(node.attrs.get("alpha", 1.0))
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.CELU,
+        shape=output_shape,
+        dtype=dtype,
+        input_dtype=dtype,
+        params=(alpha,),
+    )
+
+
+@register_lowering("Swish")
+def lower_swish(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Swish must have 1 input and 1 output")
+    dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not dtype.is_float:
+        raise UnsupportedOpError("Swish only supports floating-point inputs")
+    alpha = float(node.attrs.get("alpha", 1.0))
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.SWISH,
+        shape=output_shape,
+        dtype=dtype,
+        input_dtype=dtype,
+        params=(alpha,),
+    )
+
+
+@register_lowering("Shrink")
+def lower_shrink(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Shrink must have 1 input and 1 output")
+    dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not dtype.is_float:
+        raise UnsupportedOpError("Shrink only supports floating-point inputs")
+    bias = float(node.attrs.get("bias", 0.0))
+    lambd = float(node.attrs.get("lambd", 0.5))
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.SHRINK,
+        shape=output_shape,
+        dtype=dtype,
+        input_dtype=dtype,
+        params=(bias, lambd),
+    )
+
+
+@register_lowering("IsInf")
+def lower_isinf(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("IsInf must have 1 input and 1 output")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if not input_dtype.is_float:
+        raise UnsupportedOpError("IsInf only supports floating-point inputs")
+    if output_dtype != ScalarType.BOOL:
+        raise UnsupportedOpError("IsInf output must be bool")
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.ISINF,
+        shape=output_shape,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+        params=(),
+    )
+
+
+@register_lowering("IsNaN")
+def lower_isnan(graph: Graph, node: Node) -> UnaryOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("IsNaN must have 1 input and 1 output")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if not input_dtype.is_float:
+        raise UnsupportedOpError("IsNaN only supports floating-point inputs")
+    if output_dtype != ScalarType.BOOL:
+        raise UnsupportedOpError("IsNaN output must be bool")
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=ScalarFunction.ISNAN,
+        shape=output_shape,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+        params=(),
+    )