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

emx_onnx_cgen/lowering/one_hot.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import OneHotOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _read_scalar_initializer(
+    graph: Graph, name: str, node: Node, label: str
+) -> int | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    data = np.array(initializer.data)
+    if data.size != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} {label} input must be a scalar"
+        )
+    return int(data.reshape(-1)[0])
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+def _normalize_onehot_axis(axis: int, rank: int, node: Node) -> int:
+    if axis < 0:
+        axis += rank + 1
+    if axis < 0 or axis > rank:
+        raise ShapeInferenceError(
+            f"{node.op_type} axis {axis} is out of range for rank {rank}"
+        )
+    return axis
+
+
+@register_lowering("OneHot")
+def lower_onehot(graph: Graph, node: Node) -> OneHotOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError("OneHot must have 3 inputs and 1 output")
+    indices_name, depth_name, values_name = node.inputs
+    indices_shape = value_shape(graph, indices_name, node)
+    depth_shape = value_shape(graph, depth_name, node)
+    values_shape = value_shape(graph, values_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if not _is_scalar_shape(depth_shape):
+        raise UnsupportedOpError("OneHot depth input must be a scalar")
+    if len(values_shape) != 1 or values_shape[0] != 2:
+        raise UnsupportedOpError(
+            "OneHot values input must be a 1D tensor of size 2"
+        )
+    output_rank = len(indices_shape) + 1
+    if len(output_shape) != output_rank:
+        raise ShapeInferenceError(
+            f"OneHot output rank must be {output_rank}, got {len(output_shape)}"
+        )
+    axis = _normalize_onehot_axis(
+        int(node.attrs.get("axis", -1)), len(indices_shape), node
+    )
+    depth_value = _read_scalar_initializer(graph, depth_name, node, "depth")
+    if depth_value is not None:
+        if depth_value < 0:
+            raise ShapeInferenceError("OneHot depth must be non-negative")
+        if output_shape[axis] != depth_value:
+            raise ShapeInferenceError(
+                "OneHot output depth must be "
+                f"{depth_value}, got {output_shape[axis]}"
+            )
+        depth_dim = depth_value
+    else:
+        depth_dim = output_shape[axis]
+        if depth_dim < 0:
+            raise ShapeInferenceError("OneHot output depth must be non-negative")
+    expected_output_shape = (
+        indices_shape[:axis] + (depth_dim,) + indices_shape[axis:]
+    )
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "OneHot output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    indices_dtype = value_dtype(graph, indices_name, node)
+    depth_dtype = value_dtype(graph, depth_name, node)
+    values_dtype = value_dtype(graph, values_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if indices_dtype.is_bool:
+        raise UnsupportedOpError("OneHot indices must be numeric")
+    if depth_dtype.is_bool:
+        raise UnsupportedOpError("OneHot depth must be numeric")
+    if values_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "OneHot values dtype must match output dtype, "
+            f"got {values_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    return OneHotOp(
+        indices=indices_name,
+        depth=depth_name,
+        values=values_name,
+        output=node.outputs[0],
+        axis=axis,
+        indices_shape=indices_shape,
+        values_shape=values_shape,
+        output_shape=output_shape,
+        depth_dim=depth_dim,
+        dtype=values_dtype,
+        indices_dtype=indices_dtype,
+        depth_dtype=depth_dtype,
+    )

emx_onnx_cgen/lowering/quantize_linear.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..dtypes import scalar_type_from_onnx
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import normalize_axis
+from .common import optional_name, value_dtype as _value_dtype, value_shape as _value_shape
+from .registry import register_lowering
+from ..codegen.c_emitter import QuantizeLinearOp
+
+
+@dataclass(frozen=True)
+class QuantizeSpec:
+    input_shape: tuple[int, ...]
+    scale_shape: tuple[int, ...]
+    axis: int | None
+    output_dtype: ScalarType
+
+
+def _resolve_output_dtype(
+    graph: Graph, node: Node, zero_point_name: str | None
+) -> ScalarType:
+    output_attr = int(node.attrs.get("output_dtype", 0))
+    if output_attr:
+        output_dtype = _value_dtype(graph, node.outputs[0], node)
+        attr_dtype = scalar_type_from_onnx(output_attr)
+        if attr_dtype is None:
+            raise UnsupportedOpError(
+                "QuantizeLinear output_dtype must map to a supported scalar type"
+            )
+        if output_dtype != attr_dtype:
+            raise UnsupportedOpError(
+                "QuantizeLinear output_dtype must match output tensor dtype"
+            )
+        return output_dtype
+    if zero_point_name is None:
+        return ScalarType.U8
+    return _value_dtype(graph, zero_point_name, node)
+
+
+def resolve_quantize_spec(graph: Graph, node: Node) -> QuantizeSpec:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "QuantizeLinear must have 2 or 3 inputs and 1 output"
+        )
+    supported_attrs = {"axis", "block_size", "output_dtype", "precision", "saturate"}
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("QuantizeLinear has unsupported attributes")
+    block_size = int(node.attrs.get("block_size", 0))
+    if block_size != 0:
+        raise UnsupportedOpError("QuantizeLinear block_size is not supported")
+    precision = int(node.attrs.get("precision", 0))
+    if precision != 0:
+        raise UnsupportedOpError("QuantizeLinear precision is not supported")
+    saturate = int(node.attrs.get("saturate", 1))
+    if saturate != 1:
+        raise UnsupportedOpError("QuantizeLinear saturate must be 1")
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    scale_shape = _value_shape(graph, node.inputs[1], node)
+    zero_point_name = optional_name(node.inputs, 2)
+    output_dtype = _resolve_output_dtype(graph, node, zero_point_name)
+    if output_dtype not in {
+        ScalarType.U8,
+        ScalarType.I8,
+        ScalarType.U16,
+        ScalarType.I16,
+    }:
+        raise UnsupportedOpError(
+            "QuantizeLinear supports int8/uint8/int16/uint16 outputs only"
+        )
+    if zero_point_name is not None:
+        zero_point_dtype = _value_dtype(graph, zero_point_name, node)
+        if zero_point_dtype != output_dtype:
+            raise UnsupportedOpError(
+                "QuantizeLinear zero_point dtype must match output dtype"
+            )
+        zero_point_shape = _value_shape(graph, zero_point_name, node)
+        if zero_point_shape != scale_shape:
+            raise ShapeInferenceError(
+                "QuantizeLinear zero_point shape must match scale shape"
+            )
+    if scale_shape not in {(), (1,)}:
+        if len(scale_shape) != 1:
+            raise UnsupportedOpError(
+                "QuantizeLinear supports per-tensor and per-axis scales only"
+            )
+        axis = int(node.attrs.get("axis", 1))
+        axis = normalize_axis(axis, input_shape, node)
+        if scale_shape[0] != input_shape[axis]:
+            raise ShapeInferenceError(
+                "QuantizeLinear scale length must match input axis size"
+            )
+    else:
+        axis = None
+    return QuantizeSpec(
+        input_shape=input_shape,
+        scale_shape=scale_shape,
+        axis=axis,
+        output_dtype=output_dtype,
+    )
+
+
+@register_lowering("QuantizeLinear")
+def lower_quantize_linear(graph: Graph, node: Node) -> QuantizeLinearOp:
+    op_dtype = _value_dtype(graph, node.inputs[0], node)
+    scale_dtype = _value_dtype(graph, node.inputs[1], node)
+    if not op_dtype.is_float or not scale_dtype.is_float:
+        raise UnsupportedOpError(
+            "QuantizeLinear supports float16/float/double inputs only"
+        )
+    spec = resolve_quantize_spec(graph, node)
+    return QuantizeLinearOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        zero_point=optional_name(node.inputs, 2),
+        output=node.outputs[0],
+        input_shape=spec.input_shape,
+        axis=spec.axis,
+        dtype=spec.output_dtype,
+        input_dtype=op_dtype,
+        scale_dtype=scale_dtype,
+    )

emx_onnx_cgen/lowering/reduce.py

@@ -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/reshape.py

@@ -5,6 +5,7 @@ from shared.scalar_types import ScalarType
 from ..codegen.c_emitter 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/scatter_nd.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter 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,
+    )