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

emx_onnx_cgen/lowering/mean_variance_normalization.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import MeanVarianceNormalizationOp
+from ..errors import UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import ensure_output_shape_matches_input
+from .common import node_dtype, shape_product, value_shape
+from .reduce import normalize_reduce_axes
+from .registry import register_lowering
+
+
+@register_lowering("MeanVarianceNormalization")
+def lower_mean_variance_normalization(
+    graph: Graph, node: Node
+) -> MeanVarianceNormalizationOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "MeanVarianceNormalization must have 1 input and 1 output"
+        )
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "MeanVarianceNormalization supports float16, float, and double inputs only"
+        )
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    ensure_output_shape_matches_input(node, input_shape, output_shape)
+    axes_attr = node.attrs.get("axes")
+    if axes_attr is None:
+        axes = (0, 2, 3)
+    else:
+        axes = tuple(int(axis) for axis in axes_attr)
+    axes = normalize_reduce_axes(axes, input_shape, node)
+    if not axes:
+        raise UnsupportedOpError(
+            "MeanVarianceNormalization requires non-empty reduction axes"
+        )
+    non_axes = tuple(i for i in range(len(input_shape)) if i not in axes)
+    reduce_count = shape_product(tuple(input_shape[axis] for axis in axes))
+    return MeanVarianceNormalizationOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        shape=input_shape,
+        axes=axes,
+        non_axes=non_axes,
+        reduce_count=reduce_count,
+        epsilon=1e-9,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/negative_log_likelihood_loss.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import NegativeLogLikelihoodLossOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from .common import shape_product as _shape_product
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+def _find_node_by_output(graph: Graph, name: str) -> Node | None:
+    for node in graph.nodes:
+        if name in node.outputs:
+            return node
+    return None
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _resolve_target_shape(
+    input_shape: tuple[int, ...],
+    shape_values: list[int],
+    *,
+    allowzero: int,
+    node: Node,
+) -> tuple[int, ...]:
+    if allowzero not in (0, 1):
+        raise UnsupportedOpError("Reshape allowzero must be 0 or 1")
+    output_dims: list[int] = []
+    unknown_index: int | None = None
+    known_product = 1
+    for index, dim in enumerate(shape_values):
+        if dim == -1:
+            if unknown_index is not None:
+                raise ShapeInferenceError("Reshape allows only one -1 dimension")
+            unknown_index = index
+            output_dims.append(-1)
+            continue
+        if dim == 0:
+            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
+    input_product = _shape_product(input_shape)
+    if unknown_index is not None:
+        if known_product == 0 or input_product % known_product != 0:
+            raise ShapeInferenceError(
+                "Reshape cannot infer dimension from input 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"
+        )
+    return output_shape
+
+
+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
+    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)
+
+
+def _resolve_shape_from_reshape(
+    graph: Graph, name: str, node: Node
+) -> tuple[int, ...] | None:
+    reshape_node = _find_node_by_output(graph, name)
+    if reshape_node is None or reshape_node.op_type != "Reshape":
+        return None
+    if len(reshape_node.inputs) != 2 or len(reshape_node.outputs) != 1:
+        raise UnsupportedOpError("Reshape must have 2 inputs and 1 output")
+    input_shape = _value_shape(graph, reshape_node.inputs[0], node)
+    if not input_shape:
+        return None
+    allowzero = int(reshape_node.attrs.get("allowzero", 0))
+    shape_initializer = _find_initializer(graph, reshape_node.inputs[1])
+    if shape_initializer is not None:
+        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)]
+        return _resolve_target_shape(
+            input_shape,
+            shape_values,
+            allowzero=allowzero,
+            node=node,
+        )
+    shape_values = _shape_values_from_shape_node(
+        graph, reshape_node.inputs[1], node
+    )
+    if shape_values is None:
+        return None
+    return _resolve_target_shape(
+        input_shape,
+        shape_values,
+        allowzero=allowzero,
+        node=node,
+    )
+
+
+def _resolve_input_shape(
+    graph: Graph,
+    input_name: str,
+    target_shape: tuple[int, ...],
+    weight_name: str | None,
+    node: Node,
+) -> tuple[int, ...]:
+    input_shape = _value_shape(graph, input_name, node)
+    if input_shape:
+        return input_shape
+    reshaped = _resolve_shape_from_reshape(graph, input_name, node)
+    if reshaped is not None:
+        return reshaped
+    if weight_name is not None and target_shape:
+        weight_shape = _value_shape(graph, weight_name, node)
+        if len(weight_shape) != 1:
+            return input_shape
+        return (target_shape[0], weight_shape[0], *target_shape[1:])
+    return input_shape
+
+
+@register_lowering("NegativeLogLikelihoodLoss")
+def lower_negative_log_likelihood_loss(
+    graph: Graph, node: Node
+) -> NegativeLogLikelihoodLossOp:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "NegativeLogLikelihoodLoss must have 2 or 3 inputs and 1 output"
+        )
+    input_name = node.inputs[0]
+    target_name = node.inputs[1]
+    weight_name = node.inputs[2] if len(node.inputs) > 2 else None
+    input_dtype = _value_dtype(graph, input_name, node)
+    if not input_dtype.is_float:
+        raise UnsupportedOpError(
+            "NegativeLogLikelihoodLoss supports float16, float, and double inputs only"
+        )
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if output_dtype != input_dtype:
+        raise UnsupportedOpError(
+            "NegativeLogLikelihoodLoss output dtype must match input dtype"
+        )
+    target_dtype = _value_dtype(graph, target_name, node)
+    if target_dtype not in {ScalarType.I32, ScalarType.I64}:
+        raise UnsupportedOpError(
+            "NegativeLogLikelihoodLoss target must be int32 or int64"
+        )
+    weight_dtype = None
+    weight_shape: tuple[int, ...] | None = None
+    if weight_name is not None:
+        weight_dtype = _value_dtype(graph, weight_name, node)
+        if weight_dtype != input_dtype:
+            raise UnsupportedOpError(
+                "NegativeLogLikelihoodLoss weight dtype must match input dtype"
+            )
+    target_shape = _value_shape(graph, target_name, node)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    input_shape = _resolve_input_shape(
+        graph, input_name, target_shape, weight_name, node
+    )
+    if len(input_shape) < 2:
+        raise ShapeInferenceError(
+            "NegativeLogLikelihoodLoss input must be at least 2D"
+        )
+    if len(target_shape) != len(input_shape) - 1:
+        raise ShapeInferenceError(
+            "NegativeLogLikelihoodLoss target rank must be input rank - 1"
+        )
+    if input_shape[0] != target_shape[0]:
+        raise ShapeInferenceError(
+            "NegativeLogLikelihoodLoss target batch dimension must match input"
+        )
+    if input_shape[2:] != target_shape[1:]:
+        raise ShapeInferenceError(
+            "NegativeLogLikelihoodLoss target spatial dimensions must match input"
+        )
+    if weight_name is not None:
+        weight_shape = _value_shape(graph, weight_name, node)
+        if len(weight_shape) != 1 or weight_shape[0] != input_shape[1]:
+            raise ShapeInferenceError(
+                "NegativeLogLikelihoodLoss weight must have shape (C,)"
+            )
+    reduction = node.attrs.get("reduction", "mean")
+    if isinstance(reduction, bytes):
+        reduction = reduction.decode("utf-8")
+    if reduction not in {"none", "mean", "sum"}:
+        raise UnsupportedOpError(
+            "NegativeLogLikelihoodLoss reduction must be none, mean, or sum"
+        )
+    if reduction == "none":
+        if not output_shape:
+            output_shape = target_shape
+        if output_shape != target_shape:
+            raise ShapeInferenceError(
+                "NegativeLogLikelihoodLoss output must match target shape "
+                "when reduction is none"
+            )
+    else:
+        if output_shape and output_shape not in {(), (1,)}:
+            raise ShapeInferenceError(
+                "NegativeLogLikelihoodLoss output must be scalar when reduced"
+            )
+    n = input_shape[0]
+    c = input_shape[1]
+    d = _shape_product(input_shape[2:]) if len(input_shape) > 2 else 1
+    ignore_index = int(node.attrs.get("ignore_index", -1))
+    return NegativeLogLikelihoodLossOp(
+        input0=input_name,
+        target=target_name,
+        weight=weight_name,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        target_shape=target_shape,
+        output_shape=output_shape,
+        n=n,
+        c=c,
+        d=d,
+        reduction=reduction,
+        ignore_index=ignore_index,
+        input_dtype=input_dtype,
+        weight_dtype=weight_dtype,
+        weight_shape=weight_shape,
+        dtype=input_dtype,
+        target_dtype=target_dtype,
+    )

emx_onnx_cgen/lowering/pad.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import PadOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import optional_name, value_dtype, value_shape
+from ..validation import normalize_axis
+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_int_initializer(
+    graph: Graph,
+    name: str,
+    node: Node,
+    *,
+    label: str,
+) -> tuple[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(
+            f"Pad {label} input must be int64 or int32"
+        )
+    if len(initializer.type.shape) != 1:
+        raise UnsupportedOpError(f"Pad {label} input must be a 1D tensor")
+    values = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    return tuple(int(value) for value in values)
+
+
+def _read_scalar_initializer(
+    graph: Graph, name: str, node: Node, *, dtype: ScalarType
+) -> float | int | bool | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    if initializer.type.dtype != dtype:
+        raise UnsupportedOpError(
+            "Pad value input must match input dtype, "
+            f"got {initializer.type.dtype.onnx_name}"
+        )
+    values = np.array(initializer.data).reshape(-1)
+    if values.size != 1:
+        raise UnsupportedOpError("Pad value input must be a scalar")
+    return values.item()
+
+
+def _normalize_axes(
+    axes: tuple[int, ...], input_shape: tuple[int, ...], node: Node
+) -> tuple[int, ...]:
+    normalized = [normalize_axis(axis, input_shape, node) for axis in axes]
+    if len(set(normalized)) != len(normalized):
+        raise UnsupportedOpError("Pad axes must be unique")
+    return tuple(normalized)
+
+
+def _default_pad_value(dtype: ScalarType) -> float | int | bool:
+    if dtype.is_bool:
+        return False
+    if dtype.is_float:
+        return 0.0
+    return 0
+
+
+def _compute_strides(shape: tuple[int, ...]) -> tuple[int, ...]:
+    strides: list[int] = []
+    stride = 1
+    for dim in reversed(shape):
+        strides.append(stride)
+        stride *= dim
+    return tuple(reversed(strides))
+
+
+@register_lowering("Pad")
+def lower_pad(graph: Graph, node: Node) -> PadOp:
+    if not node.inputs or len(node.outputs) != 1:
+        raise UnsupportedOpError("Pad must have 1 output")
+    input_name = node.inputs[0]
+    if not input_name:
+        raise UnsupportedOpError("Pad input must be provided")
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Pad expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    mode = node.attrs.get("mode", "constant")
+    if isinstance(mode, bytes):
+        mode = mode.decode("utf-8")
+    if mode not in {"constant", "edge", "reflect", "wrap"}:
+        raise UnsupportedOpError(f"Pad mode '{mode}' is not supported")
+    pads_name = optional_name(node.inputs, 1)
+    pads_attr = node.attrs.get("pads")
+    if pads_name and pads_attr:
+        raise UnsupportedOpError("Pad pads must be provided via input or attribute")
+    pads = None
+    pads_input = None
+    pads_shape = None
+    pads_dtype = None
+    if pads_name:
+        pads = _read_int_initializer(graph, pads_name, node, label="pads")
+        if pads is None:
+            pads_shape = value_shape(graph, pads_name, node)
+            pads_dtype = value_dtype(graph, pads_name, node)
+            if pads_dtype not in {ScalarType.I64, ScalarType.I32}:
+                raise UnsupportedOpError(
+                    "Pad pads input must be int64 or int32"
+                )
+            if len(pads_shape) != 1:
+                raise UnsupportedOpError("Pad pads input must be a 1D tensor")
+            pads_input = pads_name
+    elif pads_attr is not None:
+        pads = tuple(int(value) for value in pads_attr)
+    if pads is None and pads_input is None:
+        pads = tuple(0 for _ in range(2 * len(input_shape)))
+
+    axes_name = optional_name(node.inputs, 3)
+    axes = None
+    axes_input = None
+    axes_shape = None
+    axes_dtype = None
+    if axes_name:
+        axes = _read_int_initializer(graph, axes_name, node, label="axes")
+        if axes is None:
+            axes_shape = value_shape(graph, axes_name, node)
+            axes_dtype = value_dtype(graph, axes_name, node)
+            if axes_dtype not in {ScalarType.I64, ScalarType.I32}:
+                raise UnsupportedOpError(
+                    "Pad axes input must be int64 or int32"
+                )
+            if len(axes_shape) != 1:
+                raise UnsupportedOpError("Pad axes input must be a 1D tensor")
+            if axes_shape[0] < 0:
+                raise ShapeInferenceError(
+                    "Pad axes input must have a static length"
+                )
+            axes_input = axes_name
+        else:
+            axes = _normalize_axes(axes, input_shape, node)
+
+    pads_axis_map = None
+    pads_values = None
+    pads_begin = None
+    pads_end = None
+
+    if axes_input is None and axes is None:
+        if pads is None:
+            if pads_shape is None or pads_shape[0] != 2 * len(input_shape):
+                raise ShapeInferenceError(
+                    "Pad pads must have length 2 * rank of input"
+                )
+            pads_begin = None
+            pads_end = None
+        else:
+            if len(pads) != 2 * len(input_shape):
+                raise ShapeInferenceError(
+                    "Pad pads must have length 2 * rank of input"
+                )
+            pads_begin = list(pads[: len(input_shape)])
+            pads_end = list(pads[len(input_shape) :])
+    elif axes_input is None:
+        if pads_input is not None:
+            if pads_shape is None or pads_shape[0] != 2 * len(axes):
+                raise ShapeInferenceError(
+                    "Pad pads must have length 2 * len(axes)"
+                )
+            pads_axis_map = [None] * len(input_shape)
+            for index, axis in enumerate(axes):
+                pads_axis_map[axis] = index
+        else:
+            if len(pads) != 2 * len(axes):
+                raise ShapeInferenceError(
+                    "Pad pads must have length 2 * len(axes)"
+                )
+            pads_begin = [0] * len(input_shape)
+            pads_end = [0] * len(input_shape)
+            for index, axis in enumerate(axes):
+                pads_begin[axis] = pads[index]
+                pads_end[axis] = pads[index + len(axes)]
+    else:
+        axes_len = axes_shape[0] if axes_shape is not None else 0
+        if pads_input is not None:
+            if pads_shape is None or pads_shape[0] != 2 * axes_len:
+                raise ShapeInferenceError(
+                    "Pad pads must have length 2 * len(axes)"
+                )
+        else:
+            if len(pads) != 2 * axes_len:
+                raise ShapeInferenceError(
+                    "Pad pads must have length 2 * len(axes)"
+                )
+            pads_values = pads
+
+    if pads_begin is not None and pads_end is not None:
+        if any(value < 0 for value in pads_begin + pads_end):
+            raise UnsupportedOpError("Pad pads must be non-negative")
+
+        expected_shape = tuple(
+            dim + pad_before + pad_after
+            for dim, pad_before, pad_after in zip(
+                input_shape, pads_begin, pads_end
+            )
+        )
+        if output_shape != expected_shape:
+            raise ShapeInferenceError(
+                "Pad output shape mismatch: "
+                f"expected {expected_shape}, got {output_shape}"
+            )
+    elif pads_values is not None:
+        if any(value < 0 for value in pads_values):
+            raise UnsupportedOpError("Pad pads must be non-negative")
+
+    value_name = optional_name(node.inputs, 2)
+    pad_value = None
+    value_input = None
+    value_input_shape = None
+    if value_name:
+        pad_value = _read_scalar_initializer(
+            graph, value_name, node, dtype=input_dtype
+        )
+        if pad_value is None:
+            value_input_shape = value_shape(graph, value_name, node)
+            input_value_dtype = value_dtype(graph, value_name, node)
+            if input_value_dtype != input_dtype:
+                raise UnsupportedOpError(
+                    "Pad value input must match input dtype, "
+                    f"got {input_value_dtype.onnx_name}"
+                )
+            if value_input_shape:
+                raise UnsupportedOpError("Pad value input must be a scalar")
+            value_input = value_name
+    elif "value" in node.attrs:
+        pad_value = node.attrs["value"]
+    if pad_value is None and value_input is None:
+        pad_value = _default_pad_value(input_dtype)
+
+    return PadOp(
+        input0=input_name,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        pads_begin=(
+            tuple(int(value) for value in pads_begin)
+            if pads_begin is not None
+            else None
+        ),
+        pads_end=(
+            tuple(int(value) for value in pads_end)
+            if pads_end is not None
+            else None
+        ),
+        pads_input=pads_input,
+        pads_shape=pads_shape,
+        pads_dtype=pads_dtype,
+        pads_axis_map=(
+            tuple(pads_axis_map) if pads_axis_map is not None else None
+        ),
+        pads_values=(
+            tuple(int(value) for value in pads_values)
+            if pads_values is not None
+            else None
+        ),
+        axes_input=axes_input,
+        axes_shape=axes_shape,
+        axes_dtype=axes_dtype,
+        mode=mode,
+        value=pad_value,
+        value_input=value_input,
+        value_shape=value_input_shape,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+        input_strides=_compute_strides(input_shape),
+    )

emx_onnx_cgen/lowering/range.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import RangeOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import node_dtype, value_shape
+from .registry import register_lowering
+
+
+_SUPPORTED_RANGE_DTYPES = {
+    ScalarType.F32,
+    ScalarType.F64,
+    ScalarType.I16,
+    ScalarType.I32,
+    ScalarType.I64,
+}
+
+
+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
+) -> float | 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 data.reshape(-1)[0].item()
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+@register_lowering("Range")
+def lower_range(graph: Graph, node: Node) -> RangeOp:
+    if len(node.inputs) != 3 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Range must have 3 inputs and 1 output")
+    start_shape = value_shape(graph, node.inputs[0], node)
+    limit_shape = value_shape(graph, node.inputs[1], node)
+    delta_shape = value_shape(graph, node.inputs[2], node)
+    if not (
+        _is_scalar_shape(start_shape)
+        and _is_scalar_shape(limit_shape)
+        and _is_scalar_shape(delta_shape)
+    ):
+        raise UnsupportedOpError("Range inputs must be scalars")
+    dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if dtype not in _SUPPORTED_RANGE_DTYPES:
+        raise UnsupportedOpError(
+            f"Range does not support dtype {dtype.onnx_name}"
+        )
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if len(output_shape) != 1:
+        raise ShapeInferenceError("Range output must be 1D")
+    start_value = _read_scalar_initializer(graph, node.inputs[0], node, "start")
+    limit_value = _read_scalar_initializer(graph, node.inputs[1], node, "limit")
+    delta_value = _read_scalar_initializer(graph, node.inputs[2], node, "delta")
+    if (
+        start_value is not None
+        and limit_value is not None
+        and delta_value is not None
+    ):
+        if float(delta_value) == 0.0:
+            raise UnsupportedOpError("Range delta must be non-zero")
+        raw_count = (
+            float(limit_value) - float(start_value)
+        ) / float(delta_value)
+        length = max(int(math.ceil(raw_count)), 0)
+        if length < 0:
+            raise ShapeInferenceError("Range output length must be non-negative")
+        if output_shape[0] != length:
+            raise ShapeInferenceError(
+                f"Range output length must be {length}, got {output_shape[0]}"
+            )
+    else:
+        length = output_shape[0]
+        if length < 0:
+            raise ShapeInferenceError("Range output length must be non-negative")
+    return RangeOp(
+        start=node.inputs[0],
+        limit=node.inputs[1],
+        delta=node.inputs[2],
+        output=node.outputs[0],
+        output_shape=output_shape,
+        length=length,
+        dtype=dtype,
+        input_dtype=dtype,
+    )