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

emx_onnx_cgen/lowering/reduce.py

@@ -6,7 +6,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReduceOp, ReshapeOp
+from ..ir.ops import ReduceOp, ReshapeOp
 from ..dtypes import scalar_type_from_onnx
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
@@ -261,13 +261,12 @@ def _infer_axes_from_shapes(
             if out_dim == in_dim:
                 if in_dim == 1:
                     return None
-                continue
-            if out_dim == 1 and in_dim != 1:
+            elif out_dim == 1 and in_dim != 1:
                 axes.append(axis)
-                continue
-            raise ShapeInferenceError(
-                f"{node.op_type} output shape does not match input shape"
-            )
+            else:
+                raise ShapeInferenceError(
+                    f"{node.op_type} output shape does not match input shape"
+                )
         return tuple(axes)
     if len(output_shape) > len(input_shape):
         return None

emx_onnx_cgen/lowering/registry.py

@@ -3,32 +3,51 @@ from __future__ import annotations
 from collections.abc import Callable, Mapping
 from typing import TypeVar
 
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
+from ..ir.op_base import OpBase
 from ..errors import UnsupportedOpError
 
 LoweredOp = TypeVar("LoweredOp")
 Handler = TypeVar("Handler")
 
-_LOWERING_REGISTRY: dict[str, Callable[[Graph, Node], object]] = {}
+_LOWERING_REGISTRY: dict[str, Callable[[Graph | GraphContext, Node], OpBase]] = {}
 
 
 def register_lowering(
     op_type: str,
 ) -> Callable[[Callable[[Graph, Node], LoweredOp]], Callable[[Graph, Node], LoweredOp]]:
     def decorator(
-        func: Callable[[Graph, Node], LoweredOp],
-    ) -> Callable[[Graph, Node], LoweredOp]:
+        func: Callable[[Graph | GraphContext, Node], LoweredOp],
+    ) -> Callable[[Graph | GraphContext, Node], LoweredOp]:
         _LOWERING_REGISTRY[op_type] = func
         return func
 
     return decorator
 
 
-def get_lowering(op_type: str) -> Callable[[Graph, Node], object] | None:
+def register_lowering_if_missing(
+    op_type: str,
+) -> Callable[[Callable[[Graph | GraphContext, Node], LoweredOp]], Callable[[Graph | GraphContext, Node], LoweredOp]]:
+    def decorator(
+        func: Callable[[Graph | GraphContext, Node], LoweredOp],
+    ) -> Callable[[Graph | GraphContext, Node], LoweredOp]:
+        if op_type not in _LOWERING_REGISTRY:
+            _LOWERING_REGISTRY[op_type] = func
+        return func
+
+    return decorator
+
+
+def get_lowering(
+    op_type: str,
+) -> Callable[[Graph | GraphContext, Node], OpBase] | None:
     return _LOWERING_REGISTRY.get(op_type)
 
 
-def get_lowering_registry() -> Mapping[str, Callable[[Graph, Node], object]]:
+def get_lowering_registry() -> Mapping[
+    str, Callable[[Graph | GraphContext, Node], OpBase]
+]:
     return _LOWERING_REGISTRY
 
 

emx_onnx_cgen/lowering/reshape.py

@@ -2,9 +2,10 @@ 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
 from .registry import register_lowering
 
 
@@ -37,6 +38,21 @@ def _shape_product(shape: tuple[int, ...]) -> int:
     return product
 
 
+def _reshape_mismatch_error(
+    node: Node,
+    input_shape: tuple[int, ...],
+    output_shape: tuple[int, ...],
+) -> ShapeInferenceError:
+    node_name = node.name or "<unnamed>"
+    return ShapeInferenceError(
+        "Reshape input/output element counts must match for op "
+        f"{node.op_type} (node '{node_name}'): input shape {input_shape}, "
+        f"output shape {output_shape}. "
+        "Hint: ensure the reshape target has the same number of elements as "
+        "the input."
+    )
+
+
 def _find_initializer(graph: Graph, name: str) -> Initializer | None:
     for initializer in graph.initializers:
         if initializer.name == name:
@@ -52,15 +68,190 @@ def _find_node_by_output(graph: Graph, name: str) -> Node | None:
 
 
 def _shape_values_from_shape_node(
-    graph: Graph, name: str, node: Node
-) -> list[int] | None:
-    shape_node = _find_node_by_output(graph, name)
-    if shape_node is None or shape_node.op_type != "Shape":
-        return None
+    graph: Graph, shape_node: Node, node: Node
+) -> list[int]:
     if len(shape_node.inputs) != 1 or len(shape_node.outputs) != 1:
         raise UnsupportedOpError("Shape must have 1 input and 1 output")
     source_shape = _value_shape(graph, shape_node.inputs[0], node)
-    return list(source_shape)
+    start = int(shape_node.attrs.get("start", 0))
+    end = int(shape_node.attrs.get("end", len(source_shape)))
+    if start < 0:
+        start += len(source_shape)
+    if end < 0:
+        end += len(source_shape)
+    start = max(start, 0)
+    end = min(end, len(source_shape))
+    if start > end:
+        return []
+    return list(source_shape[start:end])
+
+
+def _shape_values_from_initializer(
+    graph: Graph,
+    name: str,
+) -> list[int] | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    if initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            "Reshape expects int64 or int32 shape input, "
+            f"got {initializer.type.dtype.onnx_name}"
+        )
+    return [int(value) for value in initializer.data.reshape(-1)]
+
+
+def _shape_values_from_input(
+    graph: Graph,
+    name: str,
+    node: Node,
+    *,
+    _visited: set[str] | None = None,
+) -> list[int] | None:
+    if _visited is None:
+        _visited = set()
+    if name in _visited:
+        return None
+    _visited.add(name)
+    try:
+        shape_values = _shape_values_from_initializer(graph, name)
+        if shape_values is not None:
+            return shape_values
+        source_node = _find_node_by_output(graph, name)
+        if source_node is None:
+            return None
+        if source_node.op_type == "Shape":
+            return _shape_values_from_shape_node(graph, source_node, node)
+        if source_node.op_type == "Concat":
+            axis = int(source_node.attrs.get("axis", 0))
+            if axis != 0:
+                raise UnsupportedOpError("Reshape shape concat must use axis 0")
+            values: list[int] = []
+            for input_name in source_node.inputs:
+                input_values = _shape_values_from_input(
+                    graph,
+                    input_name,
+                    node,
+                    _visited=_visited,
+                )
+                if input_values is None:
+                    return None
+                values.extend(input_values)
+            return values
+        if source_node.op_type == "Cast":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Cast must have 1 input and 1 output")
+            return _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+        if source_node.op_type == "Unsqueeze":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Unsqueeze must have 1 input and 1 output")
+            return _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+        if source_node.op_type == "Identity":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Identity must have 1 input and 1 output")
+            return _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+        if source_node.op_type in {"Equal", "And", "Or", "Div", "Mod"}:
+            if len(source_node.inputs) != 2 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError(
+                    f"{source_node.op_type} must have 2 inputs and 1 output"
+                )
+            left = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            right = _shape_values_from_input(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            if left is None or right is None:
+                return None
+            if len(left) == 1 and len(right) != 1:
+                left = left * len(right)
+            if len(right) == 1 and len(left) != 1:
+                right = right * len(left)
+            if len(left) != len(right):
+                return None
+            if source_node.op_type == "Equal":
+                return [1 if l == r else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "And":
+                return [1 if (l and r) else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "Or":
+                return [1 if (l or r) else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "Div":
+                return [int(l / r) if r != 0 else 0 for l, r in zip(left, right)]
+            if source_node.op_type == "Mod":
+                return [l % r if r != 0 else 0 for l, r in zip(left, right)]
+        if source_node.op_type == "Not":
+            if len(source_node.inputs) != 1 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Not must have 1 input and 1 output")
+            values = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            if values is None:
+                return None
+            return [0 if value else 1 for value in values]
+        if source_node.op_type == "Where":
+            if len(source_node.inputs) != 3 or len(source_node.outputs) != 1:
+                raise UnsupportedOpError("Where must have 3 inputs and 1 output")
+            condition = _shape_values_from_input(
+                graph,
+                source_node.inputs[0],
+                node,
+                _visited=_visited,
+            )
+            if condition is None:
+                return None
+            on_true = _shape_values_from_input(
+                graph,
+                source_node.inputs[1],
+                node,
+                _visited=_visited,
+            )
+            on_false = _shape_values_from_input(
+                graph,
+                source_node.inputs[2],
+                node,
+                _visited=_visited,
+            )
+            if on_true is None or on_false is None:
+                return None
+            if len(condition) == 1:
+                condition = condition * max(len(on_true), len(on_false))
+            if len(on_true) == 1 and len(condition) != 1:
+                on_true = on_true * len(condition)
+            if len(on_false) == 1 and len(condition) != 1:
+                on_false = on_false * len(condition)
+            if not (len(condition) == len(on_true) == len(on_false)):
+                return None
+            return [
+                t if cond else f
+                for cond, t, f in zip(condition, on_true, on_false)
+            ]
+        return None
+    finally:
+        _visited.remove(name)
 
 
 def _resolve_target_shape(
@@ -82,19 +273,19 @@ def _resolve_target_shape(
                 raise ShapeInferenceError("Reshape allows only one -1 dimension")
             unknown_index = index
             output_dims.append(-1)
-            continue
-        if dim == 0:
-            contains_zero = True
-            if allowzero == 0:
-                if index >= len(input_shape):
-                    raise ShapeInferenceError(
-                        "Reshape zero dim must index into input shape"
-                    )
-                dim = input_shape[index]
-        if dim < 0:
-            raise ShapeInferenceError("Reshape dims must be >= -1")
-        output_dims.append(dim)
-        known_product *= dim
+        else:
+            if dim == 0:
+                contains_zero = True
+                if allowzero == 0:
+                    if index >= len(input_shape):
+                        raise ShapeInferenceError(
+                            "Reshape zero dim must index into input shape"
+                        )
+                    dim = input_shape[index]
+            if dim < 0:
+                raise ShapeInferenceError("Reshape dims must be >= -1")
+            output_dims.append(dim)
+            known_product *= dim
     if allowzero == 1 and contains_zero and unknown_index is not None:
         raise ShapeInferenceError(
             "Reshape allowzero cannot combine zero and -1 dimensions"
@@ -115,9 +306,7 @@ def _resolve_target_shape(
             output_dims[unknown_index] = input_product // known_product
     output_shape = tuple(output_dims)
     if _shape_product(output_shape) != input_product:
-        raise ShapeInferenceError(
-            "Reshape input and output element counts must match"
-        )
+        raise _reshape_mismatch_error(node, input_shape, output_shape)
     return output_shape
 
 
@@ -125,7 +314,7 @@ def _resolve_target_shape(
 def lower_reshape(graph: Graph, node: Node) -> ReshapeOp:
     if len(node.inputs) != 2 or len(node.outputs) != 1:
         raise UnsupportedOpError("Reshape must have 2 inputs and 1 output")
-    input_shape = _value_shape(graph, node.inputs[0], node)
+    input_shape = resolved_value_shape(graph, node.inputs[0], node)
     input_dtype = _value_dtype(graph, node.inputs[0], node)
     output_dtype = _value_dtype(graph, node.outputs[0], node)
     if input_dtype != output_dtype:
@@ -133,46 +322,29 @@ def lower_reshape(graph: Graph, node: Node) -> ReshapeOp:
             "Reshape expects matching input/output dtypes, "
             f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
         )
-    output_shape = _value_shape(graph, node.outputs[0], node)
+    output_value = graph.find_value(node.outputs[0])
+    output_shape = resolved_value_shape(graph, node.outputs[0], node)
+    output_dim_params = output_value.type.dim_params
     allowzero = int(node.attrs.get("allowzero", 0))
-    shape_initializer = _find_initializer(graph, node.inputs[1])
     resolved_shape: tuple[int, ...] | None = None
-    if shape_initializer is None:
-        shape_values = _shape_values_from_shape_node(
-            graph, node.inputs[1], node
-        )
-        if shape_values is not None:
-            resolved_shape = _resolve_target_shape(
-                input_shape,
-                shape_values,
-                allowzero=allowzero,
-                node=node,
-            )
-        else:
-            if _shape_product(output_shape) != _shape_product(input_shape):
-                raise ShapeInferenceError(
-                    "Reshape input and output element counts must match"
-                )
-    else:
-        if shape_initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
-            raise UnsupportedOpError(
-                "Reshape expects int64 or int32 shape input, "
-                f"got {shape_initializer.type.dtype.onnx_name}"
-            )
-        if len(shape_initializer.type.shape) != 1:
-            raise UnsupportedOpError("Reshape expects a 1D shape input")
-        shape_values = [int(value) for value in shape_initializer.data.reshape(-1)]
+    shape_values = _shape_values_from_input(graph, node.inputs[1], node)
+    if shape_values is not None:
         resolved_shape = _resolve_target_shape(
             input_shape,
             shape_values,
             allowzero=allowzero,
             node=node,
         )
-        if output_shape and resolved_shape != output_shape:
+        if output_shape and resolved_shape != output_shape and not any(
+            output_dim_params
+        ):
             raise ShapeInferenceError(
                 "Reshape output shape must be "
                 f"{resolved_shape}, got {output_shape}"
             )
+    else:
+        if _shape_product(output_shape) != _shape_product(input_shape):
+            raise _reshape_mismatch_error(node, input_shape, output_shape)
     if resolved_shape is not None:
         output_shape = resolved_shape
     for dim in output_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

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import ScatterNDOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+_ALLOWED_REDUCTIONS = {"none", "add", "mul", "min", "max"}
+
+
+@register_lowering("ScatterND")
+def lower_scatternd(graph: Graph, node: Node) -> ScatterNDOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError("ScatterND must have 3 inputs and 1 output")
+    data_name, indices_name, updates_name = node.inputs
+    output_name = node.outputs[0]
+    data_shape = value_shape(graph, data_name, node)
+    indices_shape = value_shape(graph, indices_name, node)
+    updates_shape = value_shape(graph, updates_name, node)
+    output_shape = value_shape(graph, output_name, node)
+    if output_shape != data_shape:
+        raise ShapeInferenceError(
+            "ScatterND output shape must match data shape, "
+            f"got {output_shape} vs {data_shape}"
+        )
+    if len(indices_shape) < 1:
+        raise ShapeInferenceError("ScatterND indices must have rank >= 1")
+    index_depth = indices_shape[-1]
+    if index_depth <= 0:
+        raise ShapeInferenceError(
+            "ScatterND indices final dimension must be >= 1"
+        )
+    if index_depth > len(data_shape):
+        raise ShapeInferenceError(
+            "ScatterND indices final dimension must be <= data rank, "
+            f"got {index_depth} vs {len(data_shape)}"
+        )
+    expected_updates_shape = indices_shape[:-1] + data_shape[index_depth:]
+    if updates_shape != expected_updates_shape:
+        raise ShapeInferenceError(
+            "ScatterND updates shape must be "
+            f"{expected_updates_shape}, got {updates_shape}"
+        )
+    data_dtype = value_dtype(graph, data_name, node)
+    updates_dtype = value_dtype(graph, updates_name, node)
+    if updates_dtype != data_dtype:
+        raise UnsupportedOpError(
+            "ScatterND updates dtype must match data dtype, "
+            f"got {updates_dtype.onnx_name} vs {data_dtype.onnx_name}"
+        )
+    indices_dtype = value_dtype(graph, indices_name, node)
+    if indices_dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            "ScatterND indices must be int32 or int64, "
+            f"got {indices_dtype.onnx_name}"
+        )
+    reduction_attr = node.attrs.get("reduction", "none")
+    if isinstance(reduction_attr, bytes):
+        reduction = reduction_attr.decode()
+    else:
+        reduction = str(reduction_attr)
+    if reduction not in _ALLOWED_REDUCTIONS:
+        raise UnsupportedOpError(
+            "ScatterND reduction must be one of "
+            f"{sorted(_ALLOWED_REDUCTIONS)}, got {reduction}"
+        )
+    return ScatterNDOp(
+        data=data_name,
+        indices=indices_name,
+        updates=updates_name,
+        output=output_name,
+        data_shape=data_shape,
+        indices_shape=indices_shape,
+        updates_shape=updates_shape,
+        output_shape=output_shape,
+        reduction=reduction,
+        dtype=data_dtype,
+        indices_dtype=indices_dtype,
+    )