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

emx_onnx_cgen/lowering/tfidf_vectorizer.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..ir.ops import TfIdfVectorizerOp
+from ..lowering.common import value_dtype, value_shape
+from .registry import register_lowering
+
+_SUPPORTED_INPUT_DTYPES = {ScalarType.I32, ScalarType.I64}
+_SUPPORTED_OUTPUT_DTYPES = {ScalarType.F32}
+_SUPPORTED_MODES = {"TF", "IDF", "TFIDF"}
+
+
+def _decode_mode(value: object) -> str:
+    if isinstance(value, bytes):
+        return value.decode()
+    return str(value)
+
+
+def _ensure_int_list(
+    values: object | None, *, name: str, node: Node
+) -> tuple[int, ...]:
+    if values is None:
+        raise UnsupportedOpError(f"{node.op_type} requires {name} attribute")
+    try:
+        return tuple(int(value) for value in values)  # type: ignore[arg-type]
+    except TypeError as exc:
+        raise UnsupportedOpError(
+            f"{node.op_type} {name} attribute must be a list of integers"
+        ) from exc
+
+
+def _ensure_float_list(
+    values: object | None, *, name: str, node: Node
+) -> tuple[float, ...] | None:
+    if values is None:
+        return None
+    try:
+        return tuple(float(value) for value in values)  # type: ignore[arg-type]
+    except TypeError as exc:
+        raise UnsupportedOpError(
+            f"{node.op_type} {name} attribute must be a list of floats"
+        ) from exc
+
+
+def _validate_output_shape(
+    node: Node,
+    input_shape: tuple[int, ...],
+    output_shape: tuple[int, ...],
+    output_dim: int,
+) -> None:
+    if len(input_shape) == 1:
+        expected = (output_dim,)
+    else:
+        expected = (input_shape[0], output_dim)
+    if output_shape != expected:
+        raise ShapeInferenceError(
+            f"{node.op_type} output shape must be {expected}, got {output_shape}"
+        )
+
+
+@register_lowering("TfIdfVectorizer")
+def lower_tfidf_vectorizer(graph: Graph, node: Node) -> TfIdfVectorizerOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} expects 1 input and 1 output"
+        )
+    input_name = node.inputs[0]
+    output_name = node.outputs[0]
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, output_name, node)
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, output_name, node)
+    if input_dtype not in _SUPPORTED_INPUT_DTYPES:
+        raise UnsupportedOpError(
+            f"{node.op_type} input dtype must be int32 or int64, "
+            f"got {input_dtype.onnx_name}"
+        )
+    if output_dtype not in _SUPPORTED_OUTPUT_DTYPES:
+        raise UnsupportedOpError(
+            f"{node.op_type} output dtype must be float, "
+            f"got {output_dtype.onnx_name}"
+        )
+    if len(input_shape) not in {1, 2}:
+        raise UnsupportedOpError(
+            f"{node.op_type} input rank must be 1 or 2, got {len(input_shape)}"
+        )
+    mode_value = node.attrs.get("mode")
+    if mode_value is None:
+        raise UnsupportedOpError(
+            f"{node.op_type} requires mode attribute"
+        )
+    mode = _decode_mode(mode_value)
+    if mode not in _SUPPORTED_MODES:
+        raise UnsupportedOpError(
+            f"{node.op_type} mode must be one of {sorted(_SUPPORTED_MODES)}, "
+            f"got {mode}"
+        )
+    min_gram_length = int(node.attrs.get("min_gram_length", 0))
+    max_gram_length = int(node.attrs.get("max_gram_length", 0))
+    max_skip_count = int(node.attrs.get("max_skip_count", 0))
+    if min_gram_length <= 0 or max_gram_length <= 0:
+        raise UnsupportedOpError(
+            f"{node.op_type} requires positive min/max gram lengths"
+        )
+    if min_gram_length > max_gram_length:
+        raise UnsupportedOpError(
+            f"{node.op_type} min_gram_length {min_gram_length} exceeds "
+            f"max_gram_length {max_gram_length}"
+        )
+    if max_skip_count < 0:
+        raise UnsupportedOpError(
+            f"{node.op_type} max_skip_count must be non-negative"
+        )
+    ngram_counts = _ensure_int_list(
+        node.attrs.get("ngram_counts"), name="ngram_counts", node=node
+    )
+    ngram_indexes = _ensure_int_list(
+        node.attrs.get("ngram_indexes"), name="ngram_indexes", node=node
+    )
+    if "pool_strings" in node.attrs:
+        raise UnsupportedOpError(
+            f"{node.op_type} string pools are not supported"
+        )
+    pool_int64s = _ensure_int_list(
+        node.attrs.get("pool_int64s"), name="pool_int64s", node=node
+    )
+    weights = _ensure_float_list(
+        node.attrs.get("weights"), name="weights", node=node
+    )
+    if len(ngram_counts) < max_gram_length:
+        raise UnsupportedOpError(
+            f"{node.op_type} ngram_counts length must be >= max_gram_length"
+        )
+    if ngram_counts and ngram_counts[0] != 0:
+        raise UnsupportedOpError(
+            f"{node.op_type} ngram_counts must start with 0"
+        )
+    if any(value < 0 for value in ngram_counts):
+        raise UnsupportedOpError(
+            f"{node.op_type} ngram_counts must be non-negative"
+        )
+    if any(
+        later < earlier
+        for earlier, later in zip(ngram_counts, ngram_counts[1:])
+    ):
+        raise UnsupportedOpError(
+            f"{node.op_type} ngram_counts must be non-decreasing"
+        )
+    pool_size = len(pool_int64s)
+    if ngram_counts and ngram_counts[-1] > pool_size:
+        raise UnsupportedOpError(
+            f"{node.op_type} ngram_counts exceeds pool_int64s length"
+        )
+    total_ngrams = 0
+    for gram_length in range(1, max_gram_length + 1):
+        start = ngram_counts[gram_length - 1]
+        end = (
+            ngram_counts[gram_length]
+            if gram_length < len(ngram_counts)
+            else pool_size
+        )
+        count = end - start
+        if count < 0 or count % gram_length != 0:
+            raise UnsupportedOpError(
+                f"{node.op_type} pool size for {gram_length}-grams "
+                "must be divisible by gram length"
+            )
+        total_ngrams += count // gram_length
+    if total_ngrams != len(ngram_indexes):
+        raise UnsupportedOpError(
+            f"{node.op_type} ngram_indexes length {len(ngram_indexes)} "
+            f"does not match pool ngram count {total_ngrams}"
+        )
+    if weights is not None and len(weights) != len(ngram_indexes):
+        raise UnsupportedOpError(
+            f"{node.op_type} weights length {len(weights)} does not match "
+            f"ngram_indexes length {len(ngram_indexes)}"
+        )
+    output_dim = max(ngram_indexes, default=-1) + 1
+    _validate_output_shape(node, input_shape, output_shape, output_dim)
+    return TfIdfVectorizerOp(
+        input0=input_name,
+        output=output_name,
+        input_shape=input_shape,
+        output_shape=output_shape,
+        input_dtype=input_dtype,
+        output_dtype=output_dtype,
+        min_gram_length=min_gram_length,
+        max_gram_length=max_gram_length,
+        max_skip_count=max_skip_count,
+        mode=mode,
+        ngram_counts=ngram_counts,
+        ngram_indexes=ngram_indexes,
+        pool_int64s=pool_int64s,
+        weights=weights,
+    )

emx_onnx_cgen/lowering/tile.py

@@ -30,6 +30,37 @@ def _read_repeats(graph: Graph, name: str, node: Node) -> tuple[int, ...] | None
     return tuple(int(value) for value in values)
 
 
+def _infer_repeats_from_shapes(
+    input_shape: tuple[int, ...],
+    output_shape: tuple[int, ...],
+) -> tuple[int, ...]:
+    if len(input_shape) != len(output_shape):
+        raise ShapeInferenceError(
+            "Tile repeats must have the same rank as input shape"
+        )
+    repeats: list[int] = []
+    for input_dim, output_dim in zip(input_shape, output_shape):
+        if input_dim < 0 or output_dim < 0:
+            raise ShapeInferenceError(
+                "Tile repeats input must be constant when shapes are dynamic"
+            )
+        if input_dim == 0:
+            if output_dim != 0:
+                raise ShapeInferenceError(
+                    "Tile output shape mismatch: "
+                    f"expected 0 for dimension, got {output_dim}"
+                )
+            repeats.append(0)
+            continue
+        if output_dim % input_dim != 0:
+            raise ShapeInferenceError(
+                "Tile output shape mismatch: "
+                f"expected multiple of {input_dim}, got {output_dim}"
+            )
+        repeats.append(int(output_dim // input_dim))
+    return tuple(repeats)
+
+
 def _compute_strides(shape: tuple[int, ...]) -> tuple[int, ...]:
     strides: list[int] = []
     stride = 1
@@ -54,7 +85,13 @@ def lower_tile(graph: Graph, node: Node) -> TileOp:
         )
     repeats = _read_repeats(graph, node.inputs[1], node)
     if repeats is None:
-        raise UnsupportedOpError("Tile repeats input must be a constant initializer")
+        repeats_shape = value_shape(graph, node.inputs[1], node)
+        repeats_dtype = value_dtype(graph, node.inputs[1], node)
+        if repeats_dtype not in {ScalarType.I64, ScalarType.I32}:
+            raise UnsupportedOpError("Tile repeats input must be int64 or int32")
+        if len(repeats_shape) != 1:
+            raise UnsupportedOpError("Tile repeats input must be a 1D tensor")
+        repeats = _infer_repeats_from_shapes(input_shape, output_shape)
     if len(repeats) != len(input_shape):
         raise ShapeInferenceError(
             "Tile repeats must have the same rank as input shape"

emx_onnx_cgen/lowering/topk.py

@@ -117,17 +117,13 @@ def lower_topk(graph: Graph, node: Node) -> TopKOp:
     sorted_output = bool(int(node.attrs.get("sorted", 1)))
     return TopKOp(
         input0=input_name,
+        k_input=k_name,
         output_values=output_values,
         output_indices=output_indices,
-        input_shape=input_shape,
-        output_shape=output_shape,
         axis=axis,
         k=k,
         largest=largest,
         sorted=sorted_output,
-        input_dtype=input_dtype,
-        output_values_dtype=values_dtype,
-        output_indices_dtype=indices_dtype,
     )
 
 

emx_onnx_cgen/lowering/transpose.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
-from ..ir.ops import TransposeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
+from ..ir.ops import TransposeOp
 from .common import node_dtype as _node_dtype
+from .common import value_has_dim_params as _value_has_dim_params
 from .common import value_shape as _value_shape
 from .registry import register_lowering
 
@@ -14,6 +16,8 @@ def lower_transpose(graph: Graph, node: Node) -> TransposeOp:
         raise UnsupportedOpError("Transpose 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 _value_has_dim_params(graph, node.outputs[0]) or not output_shape:
+        output_shape = ()
     perm = node.attrs.get("perm")
     if perm is None:
         perm = tuple(reversed(range(len(input_shape))))
@@ -29,18 +33,20 @@ def lower_transpose(graph: Graph, node: Node) -> TransposeOp:
             f"Transpose perm must be a permutation, got {perm}"
         )
     expected_shape = tuple(input_shape[axis] for axis in perm)
-    if output_shape != expected_shape:
+    if output_shape and output_shape != expected_shape:
         raise ShapeInferenceError(
             "Transpose output shape must match permuted input shape, "
             f"expected {expected_shape}, got {output_shape}"
         )
+    if isinstance(graph, GraphContext):
+        graph.set_shape(node.outputs[0], expected_shape)
     op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
     return TransposeOp(
         input0=node.inputs[0],
         output=node.outputs[0],
         perm=perm,
         input_shape=input_shape,
-        output_shape=output_shape,
+        output_shape=expected_shape,
         dtype=op_dtype,
         input_dtype=op_dtype,
     )

emx_onnx_cgen/lowering/unsqueeze.py

@@ -2,30 +2,20 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Initializer, Node
+from ..ir.ops import ReshapeOp
+from ..lowering.common import value_dtype, value_has_dim_params, value_shape
 from .registry import register_lowering
 
 
 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
+    return value_shape(graph, name, node)
 
 
 def _value_dtype(graph: Graph, name: str, node: Node) -> ScalarType:
-    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
+    return value_dtype(graph, name, node)
 
 
 def _find_initializer(graph: Graph, name: str) -> Initializer | None:
@@ -105,6 +95,8 @@ def lower_unsqueeze(graph: Graph, node: Node) -> ReshapeOp:
         raise UnsupportedOpError("Unsqueeze must have 1 or 2 inputs and 1 output")
     input_shape = _value_shape(graph, node.inputs[0], node)
     output_shape = _value_shape(graph, node.outputs[0], node)
+    if value_has_dim_params(graph, node.outputs[0]):
+        output_shape = ()
     _validate_shape(input_shape, node, "input")
     _validate_shape(output_shape, node, "output")
     input_dtype = _value_dtype(graph, node.inputs[0], node)
@@ -142,11 +134,14 @@ def lower_unsqueeze(graph: Graph, node: Node) -> ReshapeOp:
             )
     else:
         expected_shape = _expected_output_shape(input_shape, axes, node)
-        if expected_shape != output_shape:
+        if output_shape and expected_shape != output_shape:
             raise ShapeInferenceError(
                 "Unsqueeze output shape must be "
                 f"{expected_shape}, got {output_shape}"
             )
+        output_shape = expected_shape
+    if isinstance(graph, GraphContext):
+        graph.set_shape(node.outputs[0], output_shape)
     return ReshapeOp(
         input0=node.inputs[0],
         output=node.outputs[0],

emx_onnx_cgen/lowering/upsample.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..ir.ops import ResizeOp
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+_SUPPORTED_MODES = {"nearest", "linear"}
+
+
+def _decode_attr(value: object, default: str) -> str:
+    if value is None:
+        return default
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="ignore")
+    if isinstance(value, str):
+        return value
+    return str(value)
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _load_initializer_values(
+    graph: Graph, name: str, node: Node
+) -> tuple[float | int, ...] | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    if initializer.type.dtype not in {
+        ScalarType.F16,
+        ScalarType.F32,
+        ScalarType.F64,
+    }:
+        raise UnsupportedOpError(
+            "Upsample scales initializer must be float16/float32/float64"
+        )
+    data = initializer.data.reshape(-1)
+    return tuple(data.tolist())
+
+
+def _validate_output_shape(
+    expected: tuple[int, ...],
+    actual: tuple[int, ...],
+) -> None:
+    if expected != actual:
+        raise ShapeInferenceError(
+            f"Upsample output shape must be {expected}, got {actual}"
+        )
+    if any(dim < 0 for dim in actual):
+        raise ShapeInferenceError("Upsample output shape must be non-negative")
+
+
+@register_lowering("Upsample")
+def lower_upsample(graph: Graph, node: Node) -> ResizeOp:
+    if len(node.outputs) != 1:
+        raise UnsupportedOpError("Upsample expects one output")
+    if len(node.inputs) not in {1, 2}:
+        raise UnsupportedOpError("Upsample expects 1 or 2 inputs")
+    mode = _decode_attr(node.attrs.get("mode"), "nearest")
+    if mode not in _SUPPORTED_MODES:
+        raise UnsupportedOpError(f"Upsample mode {mode!r} is not supported")
+    input_name = node.inputs[0]
+    output_name = node.outputs[0]
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, output_name, node)
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, output_name, node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Upsample expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    rank = len(input_shape)
+    axes = tuple(range(rank))
+    scales_input = None
+    scales_shape = None
+    scales_dtype = None
+    scales_axes = None
+    scales: tuple[float, ...]
+    if len(node.inputs) == 2 and node.inputs[1]:
+        scales_input = node.inputs[1]
+        scales_shape = value_shape(graph, scales_input, node)
+        if len(scales_shape) != 1:
+            raise UnsupportedOpError("Upsample expects scales to be 1D")
+        if scales_shape[0] != rank:
+            raise UnsupportedOpError("Upsample scales length mismatch")
+        scales_dtype = value_dtype(graph, scales_input, node)
+        if scales_dtype not in {ScalarType.F16, ScalarType.F32, ScalarType.F64}:
+            raise UnsupportedOpError(
+                "Upsample expects scales input to be float16/float32/float64"
+            )
+        values = _load_initializer_values(graph, scales_input, node)
+        if values is None:
+            scales = tuple(
+                output_shape[axis] / input_shape[axis]
+                for axis in range(rank)
+            )
+        else:
+            scales = tuple(float(value) for value in values)
+            expected = tuple(
+                int(input_shape[axis] * scales[axis]) for axis in range(rank)
+            )
+            _validate_output_shape(expected, output_shape)
+    else:
+        scales_attr = node.attrs.get("scales")
+        if scales_attr is None:
+            raise UnsupportedOpError("Upsample requires scales attribute or input")
+        scales = tuple(float(value) for value in scales_attr)
+        if len(scales) != rank:
+            raise UnsupportedOpError("Upsample scales length mismatch")
+        expected = tuple(
+            int(input_shape[axis] * scales[axis]) for axis in range(rank)
+        )
+        _validate_output_shape(expected, output_shape)
+    return ResizeOp(
+        input0=input_name,
+        output=output_name,
+        input_shape=input_shape,
+        output_shape=output_shape,
+        scales=scales,
+        scales_input=scales_input,
+        sizes_input=None,
+        roi_input=None,
+        axes=axes,
+        scales_shape=scales_shape,
+        sizes_shape=None,
+        roi_shape=None,
+        scales_dtype=scales_dtype,
+        sizes_dtype=None,
+        roi_dtype=None,
+        scales_axes=scales_axes,
+        sizes_axes=None,
+        roi_axes=None,
+        mode=mode,
+        coordinate_transformation_mode="asymmetric",
+        nearest_mode="floor",
+        cubic_coeff_a=-0.75,
+        exclude_outside=False,
+        extrapolation_value=0.0,
+        antialias=False,
+        keep_aspect_ratio_policy="stretch",
+        dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/variadic.py

@@ -53,7 +53,7 @@ def _lower_variadic(graph: Graph, node: Node) -> MultiInputBinaryOp:
         output=node.outputs[0],
         function=VARIADIC_OP_FUNCTIONS[node.op_type],
         operator_kind=VARIADIC_OP_OPERATOR_KINDS[node.op_type],
-        min_inputs=2,
+        min_inputs=1 if node.op_type not in BINARY_ONLY_OPS else 2,
         max_inputs=2 if node.op_type in BINARY_ONLY_OPS else None,
     )
 

emx_onnx_cgen/lowering/where.py

@@ -65,9 +65,4 @@ def lower_where(graph: Graph, node: Node) -> WhereOp:
         input_x=x_name,
         input_y=y_name,
         output=output_name,
-        condition_shape=condition_shape,
-        x_shape=x_shape,
-        y_shape=y_shape,
-        output_shape=output_shape,
-        dtype=output_dtype,
     )