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

emx_onnx_cgen/lowering/group_normalization.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import GroupNormalizationOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import ensure_output_shape_matches_input
+from .common import node_dtype, shape_product, value_shape
+from .registry import register_lowering
+
+
+@register_lowering("GroupNormalization")
+def lower_group_normalization(
+    graph: Graph, node: Node
+) -> GroupNormalizationOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "GroupNormalization must have 3 inputs and 1 output"
+        )
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "GroupNormalization 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)
+    if len(input_shape) < 3:
+        raise ShapeInferenceError(
+            "GroupNormalization expects input rank of at least 3"
+        )
+    channels = input_shape[1]
+    num_groups_attr = node.attrs.get("num_groups")
+    if num_groups_attr is None:
+        raise UnsupportedOpError("GroupNormalization requires num_groups")
+    num_groups = int(num_groups_attr)
+    if num_groups <= 0:
+        raise ShapeInferenceError("GroupNormalization num_groups must be > 0")
+    if channels % num_groups != 0:
+        raise ShapeInferenceError(
+            "GroupNormalization num_groups must divide the channel dimension"
+        )
+    scale_shape = value_shape(graph, node.inputs[1], node)
+    bias_shape = value_shape(graph, node.inputs[2], node)
+    if scale_shape != (channels,) or bias_shape != (channels,):
+        raise ShapeInferenceError(
+            "GroupNormalization scale and bias must be 1D with length C"
+        )
+    spatial_size = shape_product(input_shape[2:])
+    group_size = channels // num_groups
+    epsilon = float(node.attrs.get("epsilon", 1e-5))
+    stash_type = int(node.attrs.get("stash_type", 1))
+    if stash_type != 1:
+        raise UnsupportedOpError(
+            "GroupNormalization supports stash_type=1 only"
+        )
+    return GroupNormalizationOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        bias=node.inputs[2],
+        output=node.outputs[0],
+        shape=input_shape,
+        channels=channels,
+        num_groups=num_groups,
+        group_size=group_size,
+        spatial_size=spatial_size,
+        epsilon=epsilon,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/identity.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import IdentityOp
+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("Identity")
+def lower_identity(graph: Graph, node: Node) -> IdentityOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Identity 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)
+    input_dim_params = graph.find_value(node.inputs[0]).type.dim_params
+    output_dim_params = graph.find_value(node.outputs[0]).type.dim_params
+    resolved_shape = output_shape or input_shape
+    if input_shape and output_shape:
+        if len(input_shape) != len(output_shape):
+            raise ShapeInferenceError("Identity input and output shapes must match")
+        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")
+    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(
+            "Identity expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    return IdentityOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        shape=resolved_shape,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/instance_normalization.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import InstanceNormalizationOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import ensure_output_shape_matches_input
+from .common import node_dtype, shape_product, value_shape
+from .registry import register_lowering
+
+
+@register_lowering("InstanceNormalization")
+def lower_instance_normalization(
+    graph: Graph, node: Node
+) -> InstanceNormalizationOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "InstanceNormalization must have 3 inputs and 1 output"
+        )
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "InstanceNormalization 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)
+    if len(input_shape) < 3:
+        raise ShapeInferenceError(
+            "InstanceNormalization expects input rank of at least 3"
+        )
+    channels = input_shape[1]
+    scale_shape = value_shape(graph, node.inputs[1], node)
+    bias_shape = value_shape(graph, node.inputs[2], node)
+    if scale_shape != (channels,) or bias_shape != (channels,):
+        raise ShapeInferenceError(
+            "InstanceNormalization scale and bias must be 1D with length C"
+        )
+    spatial_size = shape_product(input_shape[2:])
+    epsilon = float(node.attrs.get("epsilon", 1e-5))
+    return InstanceNormalizationOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        bias=node.inputs[2],
+        output=node.outputs[0],
+        shape=input_shape,
+        channels=channels,
+        spatial_size=spatial_size,
+        epsilon=epsilon,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/layer_normalization.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import LayerNormalizationOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import ensure_output_shape_matches_input
+from ..validation import normalize_axis
+from .common import node_dtype, shape_product, value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _ensure_broadcastable(
+    name: str,
+    shape: tuple[int, ...],
+    normalized_shape: tuple[int, ...],
+) -> None:
+    if len(shape) != len(normalized_shape):
+        raise ShapeInferenceError(
+            f"LayerNormalization {name} rank must match normalized rank"
+        )
+    for dim, expected in zip(shape, normalized_shape):
+        if dim not in {1, expected}:
+            raise ShapeInferenceError(
+                f"LayerNormalization {name} shape {shape} must be broadcastable "
+                f"to {normalized_shape}"
+            )
+
+
+@register_lowering("LayerNormalization")
+def lower_layer_normalization(
+    graph: Graph, node: Node
+) -> LayerNormalizationOp:
+    if len(node.inputs) < 2 or len(node.inputs) > 3:
+        raise UnsupportedOpError(
+            "LayerNormalization must have 2 or 3 inputs"
+        )
+    if len(node.outputs) < 1 or len(node.outputs) > 3:
+        raise UnsupportedOpError(
+            "LayerNormalization must have 1 to 3 outputs"
+        )
+    op_dtype = node_dtype(graph, node, *node.inputs, node.outputs[0])
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "LayerNormalization 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)
+    axis = normalize_axis(int(node.attrs.get("axis", -1)), input_shape, node)
+    normalized_shape = input_shape[axis:]
+    scale_shape = value_shape(graph, node.inputs[1], node)
+    _ensure_broadcastable("scale", scale_shape, normalized_shape)
+    bias_input = node.inputs[2] if len(node.inputs) > 2 and node.inputs[2] else None
+    bias_shape = None
+    if bias_input is not None:
+        bias_shape = value_shape(graph, bias_input, node)
+        _ensure_broadcastable("bias", bias_shape, normalized_shape)
+    epsilon = float(node.attrs.get("epsilon", 1e-5))
+    stash_type = int(node.attrs.get("stash_type", 1))
+    if stash_type != 1:
+        raise UnsupportedOpError(
+            "LayerNormalization supports stash_type=1 only"
+        )
+    mean_output = node.outputs[1] if len(node.outputs) > 1 else None
+    invstd_output = node.outputs[2] if len(node.outputs) > 2 else None
+    if mean_output is not None:
+        mean_dtype = value_dtype(graph, mean_output, node)
+        if mean_dtype != op_dtype:
+            raise UnsupportedOpError(
+                "LayerNormalization expects mean output dtype to match input"
+            )
+        expected_mean_shape = input_shape[:axis] + (1,) * len(normalized_shape)
+        mean_shape = value_shape(graph, mean_output, node)
+        if mean_shape != expected_mean_shape:
+            raise ShapeInferenceError(
+                "LayerNormalization mean output shape must be "
+                f"{expected_mean_shape}, got {mean_shape}"
+            )
+    if invstd_output is not None:
+        invstd_dtype = value_dtype(graph, invstd_output, node)
+        if invstd_dtype != op_dtype:
+            raise UnsupportedOpError(
+                "LayerNormalization expects invstd output dtype to match input"
+            )
+        expected_invstd_shape = input_shape[:axis] + (1,) * len(normalized_shape)
+        invstd_shape = value_shape(graph, invstd_output, node)
+        if invstd_shape != expected_invstd_shape:
+            raise ShapeInferenceError(
+                "LayerNormalization invstd output shape must be "
+                f"{expected_invstd_shape}, got {invstd_shape}"
+            )
+    outer = shape_product(input_shape[:axis]) if axis > 0 else 1
+    inner = shape_product(normalized_shape)
+    return LayerNormalizationOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        bias=bias_input,
+        output=node.outputs[0],
+        mean_output=mean_output,
+        invstd_output=invstd_output,
+        shape=input_shape,
+        normalized_shape=normalized_shape,
+        scale_shape=scale_shape,
+        bias_shape=bias_shape,
+        outer=outer,
+        inner=inner,
+        axis=axis,
+        epsilon=epsilon,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/logsoftmax.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+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 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("LogSoftmax")
+def lower_logsoftmax(graph: Graph, node: Node) -> LogSoftmaxOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("LogSoftmax must have 1 input and 1 output")
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "LogSoftmax 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)
+    axis = _normalize_axis(
+        int(node.attrs.get("axis", -1)),
+        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 LogSoftmaxOp(
+        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/lp_normalization.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import LpNormalizationOp
+from ..errors import UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import ensure_output_shape_matches_input
+from ..validation import normalize_axis
+from .common import node_dtype, shape_product, value_shape
+from .registry import register_lowering
+
+
+@register_lowering("LpNormalization")
+def lower_lp_normalization(graph: Graph, node: Node) -> LpNormalizationOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("LpNormalization must have 1 input and 1 output")
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "LpNormalization 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)
+    axis = normalize_axis(int(node.attrs.get("axis", -1)), input_shape, node)
+    p = int(node.attrs.get("p", 2))
+    if p not in {1, 2}:
+        raise UnsupportedOpError("LpNormalization only supports p=1 or p=2")
+    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 LpNormalizationOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        shape=input_shape,
+        axis=axis,
+        p=p,
+        outer=outer,
+        axis_size=axis_size,
+        inner=inner,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/lrn.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import LrnOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class LrnSpec:
+    shape: tuple[int, ...]
+    channels: int
+    size: int
+    half: int
+    alpha: float
+    beta: float
+    bias: float
+
+
+def _value_shape(graph: Graph, name: str, node: Node) -> tuple[int, ...]:
+    try:
+        return graph.find_value(name).type.shape
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing shape for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _value_dtype(graph: Graph, name: str, node: Node) -> str:
+    try:
+        return graph.find_value(name).type.dtype
+    except KeyError as exc:
+        raise ShapeInferenceError(
+            f"Missing dtype for value '{name}' in op {node.op_type}. "
+            "Hint: run ONNX shape inference or export with static shapes."
+        ) from exc
+
+
+def _node_dtype(graph: Graph, node: Node, *names: str) -> str:
+    dtypes = {_value_dtype(graph, name, node) for name in names}
+    if len(dtypes) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} expects matching dtypes, got {', '.join(sorted(dtypes))}"
+        )
+    return next(iter(dtypes))
+
+
+def resolve_lrn_spec(graph: Graph, node: Node) -> LrnSpec:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("LRN must have 1 input and 1 output")
+    supported_attrs = {"alpha", "beta", "bias", "size"}
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("LRN has unsupported attributes")
+    size = int(node.attrs.get("size", 0))
+    if size <= 0:
+        raise UnsupportedOpError("LRN size must be a positive integer")
+    if size % 2 == 0:
+        raise UnsupportedOpError("LRN size must be odd")
+    alpha = float(node.attrs.get("alpha", 0.0001))
+    beta = float(node.attrs.get("beta", 0.75))
+    bias = float(node.attrs.get("bias", 1.0))
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) < 2:
+        raise UnsupportedOpError("LRN expects input rank of at least 2")
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    if output_shape != input_shape:
+        raise ShapeInferenceError(
+            "LRN output shape must match input shape, "
+            f"got {output_shape} for input {input_shape}"
+        )
+    return LrnSpec(
+        shape=input_shape,
+        channels=input_shape[1],
+        size=size,
+        half=size // 2,
+        alpha=alpha,
+        beta=beta,
+        bias=bias,
+    )
+
+
+@register_lowering("LRN")
+def lower_lrn(graph: Graph, node: Node) -> LrnOp:
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "LRN supports float16, float, and double inputs only"
+        )
+    spec = resolve_lrn_spec(graph, node)
+    return LrnOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        shape=spec.shape,
+        channels=spec.channels,
+        size=spec.size,
+        half=spec.half,
+        alpha=spec.alpha,
+        beta=spec.beta,
+        bias=spec.bias,
+        dtype=op_dtype,
+    )