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

emx_onnx_cgen/lowering/reshape.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReshapeOp
+from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .common import value_shape as resolved_value_shape

emx_onnx_cgen/lowering/resize.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ResizeOp
+from ..ir.ops import ResizeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/rms_normalization.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import RMSNormalizationOp
+from ..ir.ops import RMSNormalizationOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..validation import ensure_output_shape_matches_input

emx_onnx_cgen/lowering/rotary_embedding.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import RotaryEmbeddingOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import optional_name, value_dtype, value_shape
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class RotaryEmbeddingSpec:
+    batch: int
+    seq_len: int
+    num_heads: int
+    head_size: int
+    rotary_dim: int
+    rotary_dim_half: int
+    input_rank: int
+
+
+def _resolve_rotary_spec(
+    graph: Graph, node: Node, dtype: ScalarType
+) -> RotaryEmbeddingSpec:
+    if not dtype.is_float:
+        raise UnsupportedOpError("Unsupported op RotaryEmbedding")
+    if len(node.inputs) < 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Unsupported op RotaryEmbedding")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    input_rank = len(input_shape)
+    if input_rank not in {3, 4}:
+        raise ShapeInferenceError("RotaryEmbedding expects 3D or 4D input")
+    if input_rank == 3:
+        num_heads_attr = node.attrs.get("num_heads")
+        if num_heads_attr is None:
+            raise UnsupportedOpError(
+                "RotaryEmbedding num_heads attribute is required for 3D inputs"
+            )
+        num_heads = int(num_heads_attr)
+        if num_heads <= 0:
+            raise ShapeInferenceError("RotaryEmbedding num_heads must be > 0")
+        batch, seq_len, hidden_size = input_shape
+        if hidden_size % num_heads != 0:
+            raise ShapeInferenceError(
+                "RotaryEmbedding hidden size must be divisible by num_heads"
+            )
+        head_size = hidden_size // num_heads
+    else:
+        batch, num_heads, seq_len, head_size = input_shape
+        num_heads_attr = node.attrs.get("num_heads")
+        if num_heads_attr is not None and int(num_heads_attr) != num_heads:
+            raise ShapeInferenceError(
+                "RotaryEmbedding num_heads must match input head dimension"
+            )
+    if head_size % 2 != 0:
+        raise ShapeInferenceError("RotaryEmbedding head size must be even")
+    rotary_dim = int(node.attrs.get("rotary_embedding_dim", 0))
+    if rotary_dim == 0:
+        rotary_dim = head_size
+    if rotary_dim < 0 or rotary_dim > head_size:
+        raise ShapeInferenceError(
+            "RotaryEmbedding rotary_embedding_dim must be in [0, head_size]"
+        )
+    if rotary_dim % 2 != 0:
+        raise ShapeInferenceError(
+            "RotaryEmbedding rotary_embedding_dim must be even"
+        )
+    rotary_dim_half = rotary_dim // 2
+    return RotaryEmbeddingSpec(
+        batch=batch,
+        seq_len=seq_len,
+        num_heads=num_heads,
+        head_size=head_size,
+        rotary_dim=rotary_dim,
+        rotary_dim_half=rotary_dim_half,
+        input_rank=input_rank,
+    )
+
+
+@register_lowering("RotaryEmbedding")
+def lower_rotary_embedding(graph: Graph, node: Node) -> RotaryEmbeddingOp:
+    input_name = node.inputs[0]
+    cos_name = node.inputs[1]
+    sin_name = node.inputs[2]
+    position_ids = optional_name(node.inputs, 3)
+    dtype = value_dtype(graph, input_name, node)
+    cos_dtype = value_dtype(graph, cos_name, node)
+    sin_dtype = value_dtype(graph, sin_name, node)
+    if cos_dtype != dtype or sin_dtype != dtype:
+        raise ShapeInferenceError(
+            "RotaryEmbedding inputs must share the same dtype"
+        )
+    spec = _resolve_rotary_spec(graph, node, dtype)
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if output_shape != input_shape:
+        raise ShapeInferenceError(
+            "RotaryEmbedding output shape must match input shape"
+        )
+    cos_shape = value_shape(graph, cos_name, node)
+    sin_shape = value_shape(graph, sin_name, node)
+    if cos_shape != sin_shape:
+        raise ShapeInferenceError(
+            "RotaryEmbedding cos/sin cache shapes must match"
+        )
+    position_shape = None
+    position_dtype = None
+    if position_ids is not None:
+        position_shape = value_shape(graph, position_ids, node)
+        if position_shape != (spec.batch, spec.seq_len):
+            raise ShapeInferenceError(
+                "RotaryEmbedding position_ids must match [batch, seq_len]"
+            )
+        position_dtype = value_dtype(graph, position_ids, node)
+        if not position_dtype.is_integer:
+            raise ShapeInferenceError(
+                "RotaryEmbedding position_ids must be an integer tensor"
+            )
+        if len(cos_shape) != 2:
+            raise ShapeInferenceError(
+                "RotaryEmbedding expects 2D sin/cos caches with position_ids"
+            )
+        if cos_shape[1] != spec.rotary_dim_half:
+            raise ShapeInferenceError(
+                "RotaryEmbedding cos/sin cache last dim must match rotary_dim/2"
+            )
+    else:
+        if len(cos_shape) != 3:
+            raise ShapeInferenceError(
+                "RotaryEmbedding expects 3D sin/cos caches without position_ids"
+            )
+        if cos_shape != (
+            spec.batch,
+            spec.seq_len,
+            spec.rotary_dim_half,
+        ):
+            raise ShapeInferenceError(
+                "RotaryEmbedding sin/cos cache shape must be "
+                "[batch, seq_len, rotary_dim/2]"
+            )
+    interleaved = bool(int(node.attrs.get("interleaved", 0)))
+    return RotaryEmbeddingOp(
+        input0=input_name,
+        cos_cache=cos_name,
+        sin_cache=sin_name,
+        position_ids=position_ids,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        cos_shape=cos_shape,
+        sin_shape=sin_shape,
+        position_ids_shape=position_shape,
+        dtype=dtype,
+        position_ids_dtype=position_dtype,
+        rotary_dim=spec.rotary_dim,
+        rotary_dim_half=spec.rotary_dim_half,
+        head_size=spec.head_size,
+        num_heads=spec.num_heads,
+        seq_len=spec.seq_len,
+        batch=spec.batch,
+        input_rank=spec.input_rank,
+        interleaved=interleaved,
+    )

emx_onnx_cgen/lowering/scatter_nd.py

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

emx_onnx_cgen/lowering/shape.py

@@ -2,32 +2,13 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ShapeOp
+from ..ir.ops import ShapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
+from .common import value_dtype, 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
-
-
-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
-
-
 def _normalize_slice_bounds(
     rank: int, *, start: int | None, end: int | None
 ) -> tuple[int, int]:
@@ -46,14 +27,14 @@ def _normalize_slice_bounds(
 def lower_shape(graph: Graph, node: Node) -> ShapeOp:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
         raise UnsupportedOpError("Shape 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_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
     if len(output_shape) != 1:
         raise ShapeInferenceError("Shape output must be 1D")
     if output_shape[0] < 0:
         raise ShapeInferenceError("Shape output length must be non-negative")
-    input_dtype = _value_dtype(graph, node.inputs[0], node)
-    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
     if output_dtype != ScalarType.I64:
         raise UnsupportedOpError("Shape output dtype must be int64")
     start = node.attrs.get("start")

emx_onnx_cgen/lowering/size.py

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

emx_onnx_cgen/lowering/slice.py

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

emx_onnx_cgen/lowering/softmax.py

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

emx_onnx_cgen/lowering/softmax_cross_entropy_loss.py

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

emx_onnx_cgen/lowering/split.py

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

emx_onnx_cgen/lowering/squeeze.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReshapeOp
+from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/tensor_scatter.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import TensorScatterOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import normalize_axis
+from .common import optional_name, value_dtype, value_shape
+from .registry import register_lowering
+
+_ALLOWED_MODES = {"linear", "circular"}
+
+
+@register_lowering("TensorScatter")
+def lower_tensor_scatter(graph: Graph, node: Node) -> TensorScatterOp:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "TensorScatter must have 2 or 3 inputs and 1 output"
+        )
+    past_cache_name = node.inputs[0]
+    update_name = node.inputs[1]
+    write_indices_name = optional_name(node.inputs, 2)
+    output_name = node.outputs[0]
+    past_cache_shape = value_shape(graph, past_cache_name, node)
+    update_shape = value_shape(graph, update_name, node)
+    output_shape = value_shape(graph, output_name, node)
+    if output_shape != past_cache_shape:
+        raise ShapeInferenceError(
+            "TensorScatter output shape must match past_cache shape, "
+            f"got {output_shape} vs {past_cache_shape}"
+        )
+    if len(update_shape) != len(past_cache_shape):
+        raise ShapeInferenceError(
+            "TensorScatter update shape rank must match past_cache rank, "
+            f"got {len(update_shape)} vs {len(past_cache_shape)}"
+        )
+    axis = normalize_axis(int(node.attrs.get("axis", -2)), past_cache_shape, node)
+    if axis == 0:
+        raise UnsupportedOpError(
+            "TensorScatter axis cannot be 0 (batch dimension)"
+        )
+    for dim_index, (past_dim, update_dim) in enumerate(
+        zip(past_cache_shape, update_shape)
+    ):
+        if dim_index == axis:
+            if update_dim > past_dim:
+                raise ShapeInferenceError(
+                    "TensorScatter update sequence length must be <= "
+                    "past_cache sequence length, "
+                    f"got {update_dim} vs {past_dim}"
+                )
+        elif update_dim != past_dim:
+            raise ShapeInferenceError(
+                "TensorScatter update shape must match past_cache shape "
+                f"outside axis {axis}, got {update_shape} vs {past_cache_shape}"
+            )
+    mode = node.attrs.get("mode", "linear")
+    if isinstance(mode, bytes):
+        mode = mode.decode("utf-8")
+    if mode not in _ALLOWED_MODES:
+        raise UnsupportedOpError(
+            "TensorScatter mode must be one of "
+            f"{sorted(_ALLOWED_MODES)}, got {mode}"
+        )
+    dtype = value_dtype(graph, past_cache_name, node)
+    update_dtype = value_dtype(graph, update_name, node)
+    output_dtype = value_dtype(graph, output_name, node)
+    if update_dtype != dtype or output_dtype != dtype:
+        raise UnsupportedOpError(
+            "TensorScatter expects past_cache, update, and output "
+            "to share the same dtype, "
+            f"got {dtype.onnx_name}, {update_dtype.onnx_name}, "
+            f"{output_dtype.onnx_name}"
+        )
+    write_indices_shape = None
+    write_indices_dtype = None
+    if write_indices_name is not None:
+        write_indices_shape = value_shape(graph, write_indices_name, node)
+        if len(write_indices_shape) != 1:
+            raise ShapeInferenceError(
+                "TensorScatter write_indices must be a 1D tensor"
+            )
+        if write_indices_shape[0] != past_cache_shape[0]:
+            raise ShapeInferenceError(
+                "TensorScatter write_indices length must match batch size, "
+                f"got {write_indices_shape[0]} vs {past_cache_shape[0]}"
+            )
+        write_indices_dtype = value_dtype(
+            graph, write_indices_name, node
+        )
+        if write_indices_dtype not in {ScalarType.I64, ScalarType.I32}:
+            raise UnsupportedOpError(
+                "TensorScatter write_indices must be int32 or int64, "
+                f"got {write_indices_dtype.onnx_name}"
+            )
+    return TensorScatterOp(
+        past_cache=past_cache_name,
+        update=update_name,
+        write_indices=write_indices_name,
+        output=output_name,
+        past_cache_shape=past_cache_shape,
+        update_shape=update_shape,
+        output_shape=output_shape,
+        write_indices_shape=write_indices_shape,
+        axis=axis,
+        mode=mode,
+        dtype=dtype,
+        write_indices_dtype=write_indices_dtype,
+    )

emx_onnx_cgen/lowering/tile.py

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

emx_onnx_cgen/lowering/topk.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import TopKOp
+from ..ir.ops import TopKOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import shape_product, value_dtype, value_shape
@@ -19,12 +19,10 @@ def _find_initializer(graph: Graph, name: str) -> Initializer | None:
     return None
 
 
-def _read_k(graph: Graph, name: str, node: Node) -> int:
+def _read_k(graph: Graph, name: str, node: Node) -> int | None:
     initializer = _find_initializer(graph, name)
     if initializer is None:
-        raise UnsupportedOpError(
-            f"{node.op_type} k input must be a constant initializer"
-        )
+        return None
     if initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
         raise UnsupportedOpError(
             f"{node.op_type} k input must be int64 or int32"
@@ -63,6 +61,28 @@ def lower_topk(graph: Graph, node: Node) -> TopKOp:
     axis = normalize_axis(axis, input_shape, node)
     k = _read_k(graph, k_name, node)
     axis_dim = input_shape[axis]
+    values_shape = value_shape(graph, output_values, node)
+    indices_shape = value_shape(graph, output_indices, node)
+    if values_shape != indices_shape:
+        raise ShapeInferenceError(
+            f"{node.op_type} values and indices output shapes must match, "
+            f"got {values_shape} and {indices_shape}"
+        )
+    if k is None:
+        k_shape = value_shape(graph, k_name, node)
+        if len(k_shape) != 1 or k_shape[0] != 1:
+            raise ShapeInferenceError(
+                f"{node.op_type} k input must be a 1-element tensor"
+            )
+        if axis >= len(values_shape):
+            raise ShapeInferenceError(
+                f"{node.op_type} axis {axis} exceeds output rank {len(values_shape)}"
+            )
+        k = values_shape[axis]
+        if k <= 0:
+            raise ShapeInferenceError(
+                f"{node.op_type} k must be a positive value, got {k}"
+            )
     if k > axis_dim:
         raise ShapeInferenceError(
             f"{node.op_type} k {k} exceeds axis dimension {axis_dim}"
@@ -70,12 +90,10 @@ def lower_topk(graph: Graph, node: Node) -> TopKOp:
     output_shape_expected = list(input_shape)
     output_shape_expected[axis] = k
     output_shape = tuple(output_shape_expected)
-    values_shape = value_shape(graph, output_values, node)
     if values_shape != output_shape:
         raise ShapeInferenceError(
             f"{node.op_type} values output shape must be {output_shape}, got {values_shape}"
         )
-    indices_shape = value_shape(graph, output_indices, node)
     if indices_shape != output_shape:
         raise ShapeInferenceError(
             f"{node.op_type} indices output shape must be {output_shape}, got {indices_shape}"

emx_onnx_cgen/lowering/transpose.py

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

emx_onnx_cgen/lowering/trilu.py

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

emx_onnx_cgen/lowering/unsqueeze.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReshapeOp
+from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .registry import register_lowering

emx_onnx_cgen/lowering/variadic.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from shared.scalar_functions import ScalarFunction
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import MultiInputBinaryOp
+from ..ir.ops import MultiInputBinaryOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..lowering.common import node_dtype, value_dtype, value_shape

emx_onnx_cgen/lowering/where.py

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