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

emx_onnx_cgen/lowering/__init__.py

@@ -1,3 +1,81 @@
+from __future__ import annotations
+
+import importlib
+
 from .registry import get_lowering, register_lowering
 
-__all__ = ["get_lowering", "register_lowering"]
+_LOWERING_MODULES = [
+    "adagrad",
+    "arg_reduce",
+    "attention",
+    "average_pool",
+    "batch_normalization",
+    "cast",
+    "concat",
+    "constant_of_shape",
+    "conv",
+    "conv_transpose",
+    "cumsum",
+    "depth_space",
+    "dropout",
+    "einsum",
+    "elementwise",
+    "expand",
+    "eye_like",
+    "flatten",
+    "gather",
+    "gather_elements",
+    "gather_nd",
+    "gemm",
+    "global_max_pool",
+    "grid_sample",
+    "group_normalization",
+    "hardmax",
+    "identity",
+    "instance_normalization",
+    "layer_normalization",
+    "logsoftmax",
+    "lp_normalization",
+    "lp_pool",
+    "lrn",
+    "lstm",
+    "matmul",
+    "maxpool",
+    "mean_variance_normalization",
+    "negative_log_likelihood_loss",
+    "non_max_suppression",
+    "nonzero",
+    "one_hot",
+    "pad",
+    "qlinear_matmul",
+    "quantize_linear",
+    "range",
+    "reduce",
+    "reshape",
+    "resize",
+    "rms_normalization",
+    "rotary_embedding",
+    "scatter_nd",
+    "shape",
+    "size",
+    "slice",
+    "softmax",
+    "softmax_cross_entropy_loss",
+    "split",
+    "squeeze",
+    "tensor_scatter",
+    "tile",
+    "topk",
+    "transpose",
+    "trilu",
+    "unsqueeze",
+    "variadic",
+    "where",
+]
+
+
+def load_lowering_registry() -> None:
+    for module_name in _LOWERING_MODULES:
+        importlib.import_module(f"{__name__}.{module_name}")
+
+__all__ = ["get_lowering", "register_lowering", "load_lowering_registry"]

emx_onnx_cgen/lowering/adagrad.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import AdagradOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+@register_lowering("Adagrad")
+def lower_adagrad(graph: Graph, node: Node) -> AdagradOp:
+    if len(node.inputs) < 5:
+        raise UnsupportedOpError("Adagrad must have at least 5 inputs")
+    if len(node.outputs) < 2:
+        raise UnsupportedOpError("Adagrad must have at least 2 outputs")
+    if (len(node.inputs) - 2) % 3 != 0:
+        raise UnsupportedOpError(
+            "Adagrad inputs must be R, T, Xs, Gs, Hs with matching counts"
+        )
+    tensor_count = (len(node.inputs) - 2) // 3
+    if len(node.outputs) != tensor_count * 2:
+        raise UnsupportedOpError(
+            "Adagrad outputs must be X_news followed by H_news"
+        )
+    rate_name = node.inputs[0]
+    timestep_name = node.inputs[1]
+    rate_shape = value_shape(graph, rate_name, node)
+    timestep_shape = value_shape(graph, timestep_name, node)
+    if not _is_scalar_shape(rate_shape):
+        raise UnsupportedOpError("Adagrad R input must be a scalar")
+    if not _is_scalar_shape(timestep_shape):
+        raise UnsupportedOpError("Adagrad T input must be a scalar")
+    rate_dtype = value_dtype(graph, rate_name, node)
+    if rate_dtype not in {ScalarType.F32, ScalarType.F64}:
+        raise UnsupportedOpError(
+            "Adagrad R input must be float or double"
+        )
+    timestep_dtype = value_dtype(graph, timestep_name, node)
+    if timestep_dtype != ScalarType.I64:
+        raise UnsupportedOpError("Adagrad T input must be int64")
+
+    inputs = node.inputs[2 : 2 + tensor_count]
+    gradients = node.inputs[2 + tensor_count : 2 + tensor_count * 2]
+    accumulators = node.inputs[2 + tensor_count * 2 : 2 + tensor_count * 3]
+    outputs = node.outputs[:tensor_count]
+    accumulator_outputs = node.outputs[tensor_count:]
+    if not inputs or not gradients or not accumulators:
+        raise UnsupportedOpError("Adagrad requires X, G, H inputs")
+    dtype = value_dtype(graph, inputs[0], node)
+    if dtype not in {ScalarType.F32, ScalarType.F64}:
+        raise UnsupportedOpError("Adagrad supports float and double tensors only")
+    if rate_dtype != dtype:
+        raise UnsupportedOpError(
+            "Adagrad R input dtype must match tensor dtype"
+        )
+    input_shapes: list[tuple[int, ...]] = []
+    output_shapes: list[tuple[int, ...]] = []
+    for index, (x_name, g_name, h_name, out_name, h_out_name) in enumerate(
+        zip(inputs, gradients, accumulators, outputs, accumulator_outputs)
+    ):
+        x_dtype = value_dtype(graph, x_name, node)
+        g_dtype = value_dtype(graph, g_name, node)
+        h_dtype = value_dtype(graph, h_name, node)
+        out_dtype = value_dtype(graph, out_name, node)
+        h_out_dtype = value_dtype(graph, h_out_name, node)
+        if {x_dtype, g_dtype, h_dtype, out_dtype, h_out_dtype} != {dtype}:
+            raise UnsupportedOpError(
+                "Adagrad inputs and outputs must share the same dtype"
+            )
+        x_shape = value_shape(graph, x_name, node)
+        g_shape = value_shape(graph, g_name, node)
+        h_shape = value_shape(graph, h_name, node)
+        out_shape = value_shape(graph, out_name, node)
+        h_out_shape = value_shape(graph, h_out_name, node)
+        if x_shape != g_shape or x_shape != h_shape:
+            raise ShapeInferenceError(
+                f"Adagrad inputs X/G/H shapes must match for tensor {index}"
+            )
+        if out_shape != x_shape or h_out_shape != x_shape:
+            raise ShapeInferenceError(
+                f"Adagrad outputs must match X shape for tensor {index}"
+            )
+        input_shapes.append(x_shape)
+        output_shapes.append(out_shape)
+
+    norm_coefficient = float(node.attrs.get("norm_coefficient", 0.0))
+    epsilon = float(node.attrs.get("epsilon", 0.0))
+    decay_factor = float(node.attrs.get("decay_factor", 0.0))
+
+    return AdagradOp(
+        rate=rate_name,
+        timestep=timestep_name,
+        inputs=tuple(inputs),
+        gradients=tuple(gradients),
+        accumulators=tuple(accumulators),
+        outputs=tuple(outputs),
+        accumulator_outputs=tuple(accumulator_outputs),
+        rate_shape=rate_shape,
+        timestep_shape=timestep_shape,
+        tensor_shapes=tuple(input_shapes),
+        output_shapes=tuple(output_shapes),
+        dtype=dtype,
+        rate_dtype=rate_dtype,
+        timestep_dtype=timestep_dtype,
+        norm_coefficient=norm_coefficient,
+        epsilon=epsilon,
+        decay_factor=decay_factor,
+    )

emx_onnx_cgen/lowering/arg_reduce.py

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

emx_onnx_cgen/lowering/attention.py

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

emx_onnx_cgen/lowering/average_pool.py

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

emx_onnx_cgen/lowering/batch_normalization.py

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

emx_onnx_cgen/lowering/cast.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import onnx
 
-from ..codegen.c_emitter import CastOp
+from ..ir.ops import CastOp
 from ..dtypes import scalar_type_from_onnx
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node

emx_onnx_cgen/lowering/common.py

@@ -5,6 +5,7 @@ from collections.abc import Sequence
 from shared.scalar_types import ScalarType
 
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Initializer, Node
 
 
@@ -14,7 +15,9 @@ def ensure_supported_dtype(dtype: ScalarType) -> ScalarType:
     return dtype
 
 
-def onnx_opset_version(graph: Graph, domain: str = "") -> int | None:
+def onnx_opset_version(graph: Graph | GraphContext, domain: str = "") -> int | None:
+    if isinstance(graph, GraphContext):
+        return graph.opset_version(domain)
     if domain in {"", "ai.onnx"}:
         domains = {"", "ai.onnx"}
     else:
@@ -25,7 +28,11 @@ def onnx_opset_version(graph: Graph, domain: str = "") -> int | None:
     return None
 
 
-def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType:
+def value_dtype(
+    graph: Graph | GraphContext, name: str, node: Node | None = None
+) -> ScalarType:
+    if isinstance(graph, GraphContext):
+        return graph.dtype(name, node)
     try:
         value = graph.find_value(name)
     except KeyError as exc:
@@ -37,31 +44,42 @@ def value_dtype(graph: Graph, name: str, node: Node | None = None) -> ScalarType
     return ensure_supported_dtype(value.type.dtype)
 
 
-def value_shape(graph: Graph, name: str, node: Node | None = None) -> tuple[int, ...]:
-    try:
+def value_shape(
+    graph: Graph | GraphContext, name: str, node: Node | None = None
+) -> tuple[int, ...]:
+    if isinstance(graph, GraphContext):
+        shape = graph.shape(name, node)
         value = graph.find_value(name)
-    except KeyError as exc:
-        op_type = node.op_type if node is not None else "unknown"
-        raise ShapeInferenceError(
-            f"Missing shape for value '{name}' in op {op_type}. "
-            "Hint: run ONNX shape inference or export with static shapes."
-        ) from exc
+    else:
+        try:
+            value = graph.find_value(name)
+        except KeyError as exc:
+            op_type = node.op_type if node is not None else "unknown"
+            raise ShapeInferenceError(
+                f"Missing shape for value '{name}' in op {op_type}. "
+                "Hint: run ONNX shape inference or export with static shapes."
+            ) from exc
+        shape = value.type.shape
     if any(value.type.dim_params):
         resolved = _resolve_value_shape(graph, name, node)
         if resolved is not None:
             return resolved
         return value.type.shape
-    return value.type.shape
+    return shape
 
 
-def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+def _find_initializer(graph: Graph | GraphContext, name: str) -> Initializer | None:
+    if isinstance(graph, GraphContext):
+        return graph.initializer(name)
     for initializer in graph.initializers:
         if initializer.name == name:
             return initializer
     return None
 
 
-def _find_node_by_output(graph: Graph, name: str) -> Node | None:
+def _find_node_by_output(graph: Graph | GraphContext, name: str) -> Node | None:
+    if isinstance(graph, GraphContext):
+        return graph.producer(name)
     for node in graph.nodes:
         if name in node.outputs:
             return node
@@ -69,7 +87,7 @@ def _find_node_by_output(graph: Graph, name: str) -> Node | None:
 
 
 def _shape_values_from_shape_node(
-    graph: Graph, shape_node: Node, node: Node | None
+    graph: Graph | GraphContext, shape_node: Node, node: Node | None
 ) -> list[int]:
     if len(shape_node.inputs) != 1 or len(shape_node.outputs) != 1:
         raise UnsupportedOpError("Shape must have 1 input and 1 output")
@@ -88,7 +106,7 @@ def _shape_values_from_shape_node(
 
 
 def _shape_values_from_initializer(
-    graph: Graph,
+    graph: Graph | GraphContext,
     name: str,
 ) -> list[int] | None:
     initializer = _find_initializer(graph, name)
@@ -103,7 +121,7 @@ def _shape_values_from_initializer(
 
 
 def _shape_values_from_input(
-    graph: Graph,
+    graph: Graph | GraphContext,
     name: str,
     node: Node | None,
     *,
@@ -277,7 +295,7 @@ def _broadcast_shapes(
 
 
 def _resolve_value_shape(
-    graph: Graph,
+    graph: Graph | GraphContext,
     name: str,
     node: Node | None,
     *,
@@ -414,7 +432,7 @@ def _resolve_value_shape(
         _visited.remove(name)
 
 
-def node_dtype(graph: Graph, node: Node, *names: str) -> ScalarType:
+def node_dtype(graph: Graph | GraphContext, node: Node, *names: str) -> ScalarType:
     filtered = [name for name in names if name]
     if not filtered:
         raise UnsupportedOpError(

emx_onnx_cgen/lowering/concat.py

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

emx_onnx_cgen/lowering/constant_of_shape.py

@@ -4,7 +4,7 @@ from onnx import numpy_helper
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ConstantOfShapeOp
+from ..ir.ops import ConstantOfShapeOp
 from ..dtypes import scalar_type_from_onnx
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node

emx_onnx_cgen/lowering/conv.py

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

emx_onnx_cgen/lowering/conv_transpose.py

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

emx_onnx_cgen/lowering/cumsum.py

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

emx_onnx_cgen/lowering/depth_space.py

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

emx_onnx_cgen/lowering/dropout.py

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

emx_onnx_cgen/lowering/einsum.py

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

emx_onnx_cgen/lowering/elementwise.py

@@ -1,13 +1,23 @@
 from __future__ import annotations
 
-from shared.scalar_functions import ScalarFunction
+from shared.scalar_functions import ScalarFunction, ScalarFunctionError
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ClipOp, UnaryOp
+from ..ir.ops import BinaryOp, ClipOp, UnaryOp
 from ..errors import UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
 from ..lowering.common import node_dtype, optional_name, value_dtype, value_shape
-from ..lowering.registry import register_lowering
+from ..lowering.registry import register_lowering, register_lowering_if_missing
+from ..ops import (
+    BINARY_OP_TYPES,
+    COMPARE_FUNCTIONS,
+    UNARY_OP_TYPES,
+    binary_op_symbol,
+    unary_op_symbol,
+    validate_unary_attrs,
+)
+from ..lowering.variadic import VARIADIC_OP_FUNCTIONS
 
 
 @register_lowering("Clip")
@@ -120,6 +130,138 @@ def lower_shrink(graph: Graph, node: Node) -> UnaryOp:
     )
 
 
+def _lower_binary_unary(graph: Graph | GraphContext, node: Node) -> BinaryOp | UnaryOp:
+    if node.op_type == "BitShift":
+        if len(node.inputs) != 2 or len(node.outputs) != 1:
+            raise UnsupportedOpError("BitShift must have 2 inputs and 1 output")
+        direction_attr = node.attrs.get("direction", "LEFT")
+        if isinstance(direction_attr, bytes):
+            direction = direction_attr.decode()
+        else:
+            direction = str(direction_attr)
+        if direction not in {"LEFT", "RIGHT"}:
+            raise UnsupportedOpError(
+                "BitShift direction must be LEFT or RIGHT"
+            )
+        op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+        if not op_dtype.is_integer:
+            raise UnsupportedOpError("BitShift expects integer inputs")
+        function = (
+            ScalarFunction.BITWISE_LEFT_SHIFT
+            if direction == "LEFT"
+            else ScalarFunction.BITWISE_RIGHT_SHIFT
+        )
+        op_spec = binary_op_symbol(function, node.attrs, dtype=op_dtype)
+        if op_spec is None:
+            raise UnsupportedOpError("Unsupported op BitShift")
+        input0_shape = value_shape(graph, node.inputs[0], node)
+        input1_shape = value_shape(graph, node.inputs[1], node)
+        output_shape = value_shape(graph, node.outputs[0], node)
+        return BinaryOp(
+            input0=node.inputs[0],
+            input1=node.inputs[1],
+            output=node.outputs[0],
+            function=function,
+            operator_kind=op_spec.kind,
+            input0_shape=input0_shape,
+            input1_shape=input1_shape,
+            shape=output_shape,
+            dtype=op_dtype,
+            input_dtype=op_dtype,
+        )
+    if node.op_type == "Mod":
+        fmod = int(node.attrs.get("fmod", 0))
+        if fmod not in {0, 1}:
+            raise UnsupportedOpError("Mod only supports fmod=0 or fmod=1")
+        function = (
+            ScalarFunction.FMOD if fmod == 1 else ScalarFunction.REMAINDER
+        )
+    else:
+        try:
+            function = ScalarFunction.from_onnx_op(node.op_type)
+        except ScalarFunctionError as exc:
+            raise UnsupportedOpError(
+                f"Unsupported op {node.op_type}"
+            ) from exc
+    validate_unary_attrs(node.op_type, node.attrs)
+    if function in COMPARE_FUNCTIONS:
+        input_dtype = node_dtype(graph, node, *node.inputs)
+        output_dtype = value_dtype(graph, node.outputs[0], node)
+        op_spec = binary_op_symbol(function, node.attrs, dtype=input_dtype)
+        if op_spec is None:
+            raise UnsupportedOpError(f"Unsupported op {node.op_type}")
+        if len(node.inputs) != 2 or len(node.outputs) != 1:
+            raise UnsupportedOpError(
+                f"{node.op_type} must have 2 inputs and 1 output"
+            )
+        if output_dtype != ScalarType.BOOL:
+            raise UnsupportedOpError(
+                f"{node.op_type} expects bool output, got {output_dtype.onnx_name}"
+            )
+        input0_shape = value_shape(graph, node.inputs[0], node)
+        input1_shape = value_shape(graph, node.inputs[1], node)
+        output_shape = value_shape(graph, node.outputs[0], node)
+        return BinaryOp(
+            input0=node.inputs[0],
+            input1=node.inputs[1],
+            output=node.outputs[0],
+            function=function,
+            operator_kind=op_spec.kind,
+            input0_shape=input0_shape,
+            input1_shape=input1_shape,
+            shape=output_shape,
+            dtype=output_dtype,
+            input_dtype=input_dtype,
+        )
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    op_spec = binary_op_symbol(function, node.attrs, dtype=op_dtype)
+    unary_symbol = unary_op_symbol(function, dtype=op_dtype)
+    if op_spec is None and unary_symbol is None:
+        raise UnsupportedOpError(f"Unsupported op {node.op_type}")
+    if op_spec is not None:
+        if len(node.inputs) != 2 or len(node.outputs) != 1:
+            raise UnsupportedOpError(
+                f"{node.op_type} must have 2 inputs and 1 output"
+            )
+        input0_shape = value_shape(graph, node.inputs[0], node)
+        input1_shape = value_shape(graph, node.inputs[1], node)
+        output_shape = value_shape(graph, node.outputs[0], node)
+        return BinaryOp(
+            input0=node.inputs[0],
+            input1=node.inputs[1],
+            output=node.outputs[0],
+            function=function,
+            operator_kind=op_spec.kind,
+            input0_shape=input0_shape,
+            input1_shape=input1_shape,
+            shape=output_shape,
+            dtype=op_dtype,
+            input_dtype=op_dtype,
+        )
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} must have 1 input and 1 output"
+        )
+    output_shape = value_shape(graph, node.outputs[0], node)
+    return UnaryOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        function=function,
+        shape=output_shape,
+        dtype=op_dtype,
+        input_dtype=op_dtype,
+        params=(),
+    )
+
+
+_DEFAULT_ELEMENTWISE_TYPES = (
+    BINARY_OP_TYPES.union(UNARY_OP_TYPES) - set(VARIADIC_OP_FUNCTIONS.keys())
+)
+
+for _op_type in _DEFAULT_ELEMENTWISE_TYPES:
+    register_lowering_if_missing(_op_type)(_lower_binary_unary)
+
+
 @register_lowering("IsInf")
 def lower_isinf(graph: Graph, node: Node) -> UnaryOp:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
@@ -130,6 +272,12 @@ def lower_isinf(graph: Graph, node: Node) -> UnaryOp:
         raise UnsupportedOpError("IsInf only supports floating-point inputs")
     if output_dtype != ScalarType.BOOL:
         raise UnsupportedOpError("IsInf output must be bool")
+    detect_negative = int(node.attrs.get("detect_negative", 1))
+    detect_positive = int(node.attrs.get("detect_positive", 1))
+    if detect_negative not in {0, 1} or detect_positive not in {0, 1}:
+        raise UnsupportedOpError(
+            "IsInf detect_negative and detect_positive must be 0 or 1"
+        )
     output_shape = value_shape(graph, node.outputs[0], node)
     return UnaryOp(
         input0=node.inputs[0],
@@ -138,7 +286,7 @@ def lower_isinf(graph: Graph, node: Node) -> UnaryOp:
         shape=output_shape,
         dtype=output_dtype,
         input_dtype=input_dtype,
-        params=(),
+        params=(float(detect_negative), float(detect_positive)),
     )
 
 

emx_onnx_cgen/lowering/expand.py

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

emx_onnx_cgen/lowering/eye_like.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import EyeLikeOp
+from ..ir.ops import EyeLikeOp
 from ..dtypes import scalar_type_from_onnx
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node

emx_onnx_cgen/lowering/flatten.py

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

emx_onnx_cgen/lowering/gather.py

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