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

emx_onnx_cgen/lowering/reshape.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+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 .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 _shape_product(shape: tuple[int, ...]) -> int:
+    product = 1
+    for dim in shape:
+        if dim < 0:
+            raise ShapeInferenceError("Dynamic dims are not supported")
+        product *= dim
+    return product
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+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 _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_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
+    contains_zero = False
+    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:
+            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"
+        )
+    input_product = _shape_product(input_shape)
+    if unknown_index is not None:
+        if known_product == 0:
+            if input_product != 0:
+                raise ShapeInferenceError(
+                    "Reshape cannot infer dimension from input shape"
+                )
+            output_dims[unknown_index] = 0
+        else:
+            if 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
+
+
+@register_lowering("Reshape")
+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_dtype = _value_dtype(graph, node.inputs[0], node)
+    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "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)
+    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)]
+        resolved_shape = _resolve_target_shape(
+            input_shape,
+            shape_values,
+            allowzero=allowzero,
+            node=node,
+        )
+        if output_shape and resolved_shape != output_shape:
+            raise ShapeInferenceError(
+                "Reshape output shape must be "
+                f"{resolved_shape}, got {output_shape}"
+            )
+    if resolved_shape is not None:
+        output_shape = resolved_shape
+    for dim in output_shape:
+        if dim < 0:
+            raise ShapeInferenceError("Dynamic dims are not supported")
+    return ReshapeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        dtype=input_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/resize.py

@@ -0,0 +1,445 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import ResizeOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from .registry import register_lowering
+
+_SUPPORTED_COORD_MODES = {
+    "half_pixel",
+    "half_pixel_symmetric",
+    "asymmetric",
+    "align_corners",
+    "pytorch_half_pixel",
+    "tf_crop_and_resize",
+}
+_SUPPORTED_MODES = {"nearest", "linear", "cubic"}
+_SUPPORTED_NEAREST_MODES = {
+    "round_prefer_floor",
+    "round_prefer_ceil",
+    "floor",
+    "ceil",
+}
+_SUPPORTED_KEEP_ASPECT = {"stretch", "not_larger", "not_smaller"}
+
+
+@dataclass(frozen=True)
+class _ResizeInputs:
+    roi: str | None
+    scales: str | None
+    sizes: str | None
+
+
+@dataclass(frozen=True)
+class _ResolvedScales:
+    scales: tuple[float, ...]
+    output_shape: tuple[int, ...]
+    axes: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class _InputConfig:
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    input_dtype: ScalarType
+    output_dtype: ScalarType
+
+
+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 _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _decode_attr(value: object, default: str) -> str:
+    if value is None:
+        return default
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="ignore")
+    if isinstance(value, str):
+        return value
+    return str(value)
+
+
+def _normalize_axes(
+    axes: tuple[int, ...], rank: int, node: Node
+) -> tuple[int, ...]:
+    normalized: list[int] = []
+    for axis in axes:
+        axis = int(axis)
+        if axis < 0:
+            axis += rank
+        if axis < 0 or axis >= rank:
+            raise ShapeInferenceError(
+                f"Resize axis {axis} is out of range for rank {rank}"
+            )
+        normalized.append(axis)
+    if len(set(normalized)) != len(normalized):
+        raise ShapeInferenceError("Resize axes must be unique")
+    return tuple(normalized)
+
+
+def _round_half_up(value: float) -> int:
+    return int(value + 0.5)
+
+
+def _parse_input_names(node: Node) -> _ResizeInputs:
+    inputs = list(node.inputs)
+    if len(inputs) > 4:
+        raise UnsupportedOpError("Resize expects at most 4 inputs")
+    while len(inputs) < 4:
+        inputs.append("")
+    _, roi, scales, sizes = inputs[:4]
+    return _ResizeInputs(
+        roi=roi or None,
+        scales=scales or None,
+        sizes=sizes or None,
+    )
+
+
+def _parse_axes(node: Node, rank: int) -> tuple[int, ...]:
+    axes_attr = node.attrs.get("axes")
+    if axes_attr is None:
+        return tuple(range(rank))
+    axes = tuple(int(value) for value in axes_attr)
+    return _normalize_axes(axes, rank, node)
+
+
+def _resolve_input_shapes(
+    graph: Graph, node: Node, input_name: str
+) -> _InputConfig:
+    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)
+    return _InputConfig(
+        input_shape=input_shape,
+        output_shape=output_shape,
+        input_dtype=input_dtype,
+        output_dtype=output_dtype,
+    )
+
+
+def _resolve_scales_from_sizes(
+    sizes: tuple[int, ...],
+    input_shape: tuple[int, ...],
+    axes: tuple[int, ...],
+    keep_aspect_ratio_policy: str,
+) -> _ResolvedScales:
+    rank = len(input_shape)
+    full_sizes = list(input_shape)
+    for index, axis in enumerate(axes):
+        full_sizes[axis] = sizes[index]
+    if keep_aspect_ratio_policy != "stretch":
+        scales = [full_sizes[axis] / input_shape[axis] for axis in axes]
+        if keep_aspect_ratio_policy == "not_larger":
+            scale = min(scales)
+        else:
+            scale = max(scales)
+        for axis in axes:
+            full_sizes[axis] = _round_half_up(scale * input_shape[axis])
+        return _ResolvedScales(
+            scales=tuple(
+                scale if axis in axes else 1.0
+                for axis in range(rank)
+            ),
+            output_shape=tuple(full_sizes),
+            axes=axes,
+        )
+    scales = tuple(
+        full_sizes[axis] / input_shape[axis] if axis in axes else 1.0
+        for axis in range(rank)
+    )
+    return _ResolvedScales(
+        scales=scales,
+        output_shape=tuple(full_sizes),
+        axes=axes,
+    )
+
+
+def _resolve_scales_from_values(
+    scales: tuple[float, ...],
+    input_shape: tuple[int, ...],
+    axes: tuple[int, ...],
+) -> _ResolvedScales:
+    rank = len(input_shape)
+    full_scales = [1.0] * rank
+    for index, axis in enumerate(axes):
+        full_scales[axis] = scales[index]
+    output_shape = tuple(
+        int(input_shape[axis] * full_scales[axis])
+        if axis in axes
+        else input_shape[axis]
+        for axis in range(rank)
+    )
+    return _ResolvedScales(
+        scales=tuple(full_scales),
+        output_shape=output_shape,
+        axes=axes,
+    )
+
+
+def _load_initializer_values(
+    graph: Graph, name: str, node: Node
+) -> tuple[float | int, ...] | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    data = initializer.data.reshape(-1)
+    return tuple(data.tolist())
+
+
+def _validate_tensor_1d(
+    graph: Graph,
+    name: str,
+    node: Node,
+    dtype_options: set[ScalarType],
+) -> tuple[int, ScalarType]:
+    shape = _value_shape(graph, name, node)
+    if len(shape) != 1:
+        raise UnsupportedOpError("Resize expects 1D auxiliary inputs")
+    dtype = _value_dtype(graph, name, node)
+    if dtype not in dtype_options:
+        raise UnsupportedOpError(
+            "Resize expects "
+            f"{name} to have dtype in {[dtype.onnx_name for dtype in sorted(dtype_options, key=str)]}"
+        )
+    return shape[0], dtype
+
+
+def _resolve_scales(
+    graph: Graph,
+    node: Node,
+    config: _InputConfig,
+    inputs: _ResizeInputs,
+    axes: tuple[int, ...],
+    keep_aspect_ratio_policy: str,
+) -> tuple[tuple[float, ...], tuple[int, ...]]:
+    rank = len(config.input_shape)
+    if inputs.scales:
+        scale_len, _ = _validate_tensor_1d(
+            graph,
+            inputs.scales,
+            node,
+            {ScalarType.F16, ScalarType.F32, ScalarType.F64},
+        )
+        if scale_len not in {len(axes), rank}:
+            raise UnsupportedOpError("Resize scales length mismatch")
+        if scale_len == rank and axes != tuple(range(rank)):
+            raise UnsupportedOpError(
+                "Resize scales length conflicts with axes configuration"
+            )
+        scale_axes = axes if scale_len == len(axes) else tuple(range(rank))
+        values = _load_initializer_values(graph, inputs.scales, node)
+        if values is None:
+            scales = tuple(
+                config.output_shape[axis] / config.input_shape[axis]
+                if axis in scale_axes
+                else 1.0
+                for axis in range(rank)
+            )
+            return scales, config.output_shape
+        resolved = _resolve_scales_from_values(
+            tuple(float(value) for value in values),
+            config.input_shape,
+            scale_axes,
+        )
+        return resolved.scales, resolved.output_shape
+    if inputs.sizes:
+        size_len, _ = _validate_tensor_1d(
+            graph, inputs.sizes, node, {ScalarType.I64, ScalarType.I32}
+        )
+        if size_len not in {len(axes), rank}:
+            raise UnsupportedOpError("Resize sizes length mismatch")
+        if size_len == rank and axes != tuple(range(rank)):
+            raise UnsupportedOpError(
+                "Resize sizes length conflicts with axes configuration"
+            )
+        size_axes = axes if size_len == len(axes) else tuple(range(rank))
+        values = _load_initializer_values(graph, inputs.sizes, node)
+        if values is None:
+            scales = tuple(
+                config.output_shape[axis] / config.input_shape[axis]
+                if axis in size_axes
+                else 1.0
+                for axis in range(rank)
+            )
+            return scales, config.output_shape
+        resolved = _resolve_scales_from_sizes(
+            tuple(int(value) for value in values),
+            config.input_shape,
+            size_axes,
+            keep_aspect_ratio_policy,
+        )
+        return resolved.scales, resolved.output_shape
+    raise UnsupportedOpError("Resize expects scales or sizes input")
+
+
+def _validate_output_shape(
+    expected: tuple[int, ...],
+    actual: tuple[int, ...],
+) -> None:
+    if expected != actual:
+        raise ShapeInferenceError(
+            f"Resize output shape must be {expected}, got {actual}"
+        )
+    if any(dim < 0 for dim in actual):
+        raise ShapeInferenceError("Resize output shape must be non-negative")
+
+
+@register_lowering("Resize")
+def lower_resize(graph: Graph, node: Node) -> ResizeOp:
+    if len(node.outputs) != 1:
+        raise UnsupportedOpError("Resize expects one output")
+    inputs = _parse_input_names(node)
+    if inputs.scales and inputs.sizes:
+        raise UnsupportedOpError("Resize cannot set both scales and sizes")
+    if not inputs.scales and not inputs.sizes:
+        raise UnsupportedOpError("Resize expects scales or sizes input")
+    mode = _decode_attr(node.attrs.get("mode"), "nearest")
+    coordinate_mode = _decode_attr(
+        node.attrs.get("coordinate_transformation_mode"), "half_pixel"
+    )
+    nearest_mode = _decode_attr(
+        node.attrs.get("nearest_mode"), "round_prefer_floor"
+    )
+    keep_aspect_ratio_policy = _decode_attr(
+        node.attrs.get("keep_aspect_ratio_policy"), "stretch"
+    )
+    antialias = bool(int(node.attrs.get("antialias", 0)))
+    cubic_coeff_a = float(node.attrs.get("cubic_coeff_a", -0.75))
+    exclude_outside = bool(int(node.attrs.get("exclude_outside", 0)))
+    extrapolation_value = float(node.attrs.get("extrapolation_value", 0.0))
+    if mode not in _SUPPORTED_MODES:
+        raise UnsupportedOpError(f"Resize mode {mode!r} is not supported")
+    if coordinate_mode not in _SUPPORTED_COORD_MODES:
+        raise UnsupportedOpError(
+            "Resize coordinate_transformation_mode "
+            f"{coordinate_mode!r} is not supported"
+        )
+    if nearest_mode not in _SUPPORTED_NEAREST_MODES:
+        raise UnsupportedOpError(
+            f"Resize nearest_mode {nearest_mode!r} is not supported"
+        )
+    if keep_aspect_ratio_policy not in _SUPPORTED_KEEP_ASPECT:
+        raise UnsupportedOpError(
+            "Resize keep_aspect_ratio_policy "
+            f"{keep_aspect_ratio_policy!r} is not supported"
+        )
+    if antialias and mode == "nearest":
+        raise UnsupportedOpError("Resize antialias is not supported for nearest")
+    config = _resolve_input_shapes(graph, node, node.inputs[0])
+    if config.input_dtype != config.output_dtype:
+        raise UnsupportedOpError(
+            "Resize expects matching input/output dtypes, "
+            f"got {config.input_dtype.onnx_name} and {config.output_dtype.onnx_name}"
+        )
+    rank = len(config.input_shape)
+    axes = _parse_axes(node, rank)
+    scales, expected_output = _resolve_scales(
+        graph,
+        node,
+        config,
+        inputs,
+        axes,
+        keep_aspect_ratio_policy,
+    )
+    _validate_output_shape(expected_output, config.output_shape)
+    roi_shape = None
+    roi_axes = None
+    roi_dtype = None
+    if inputs.roi:
+        roi_len, roi_dtype = _validate_tensor_1d(
+            graph,
+            inputs.roi,
+            node,
+            {ScalarType.F16, ScalarType.F32, ScalarType.F64},
+        )
+        if roi_len == 2 * rank:
+            roi_shape = (roi_len,)
+        elif roi_len == 2 * len(axes):
+            roi_shape = (roi_len,)
+            roi_axes = axes
+        else:
+            raise UnsupportedOpError("Resize roi length mismatch")
+        if coordinate_mode != "tf_crop_and_resize" and roi_len != 0:
+            roi_axes = roi_axes if roi_len == 2 * len(axes) else None
+    if coordinate_mode == "tf_crop_and_resize" and not inputs.roi:
+        raise UnsupportedOpError("Resize requires roi for tf_crop_and_resize")
+    scales_shape = None
+    sizes_shape = None
+    scales_dtype = None
+    sizes_dtype = None
+    scales_axes = None
+    sizes_axes = None
+    if inputs.scales:
+        scale_len, scales_dtype = _validate_tensor_1d(
+            graph,
+            inputs.scales,
+            node,
+            {ScalarType.F16, ScalarType.F32, ScalarType.F64},
+        )
+        scales_shape = (scale_len,)
+        if scale_len == len(axes) and len(axes) != rank:
+            scales_axes = axes
+    if inputs.sizes:
+        size_len, sizes_dtype = _validate_tensor_1d(
+            graph, inputs.sizes, node, {ScalarType.I64, ScalarType.I32}
+        )
+        sizes_shape = (size_len,)
+        if size_len == len(axes) and len(axes) != rank:
+            sizes_axes = axes
+    return ResizeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=config.input_shape,
+        output_shape=config.output_shape,
+        scales=scales,
+        scales_input=inputs.scales,
+        sizes_input=inputs.sizes,
+        roi_input=inputs.roi,
+        axes=axes,
+        scales_shape=scales_shape,
+        sizes_shape=sizes_shape,
+        roi_shape=roi_shape,
+        scales_dtype=scales_dtype,
+        sizes_dtype=sizes_dtype,
+        roi_dtype=roi_dtype,
+        scales_axes=scales_axes,
+        sizes_axes=sizes_axes,
+        roi_axes=roi_axes,
+        mode=mode,
+        coordinate_transformation_mode=coordinate_mode,
+        nearest_mode=nearest_mode,
+        cubic_coeff_a=cubic_coeff_a,
+        exclude_outside=exclude_outside,
+        extrapolation_value=extrapolation_value,
+        antialias=antialias,
+        keep_aspect_ratio_policy=keep_aspect_ratio_policy,
+        dtype=config.input_dtype,
+    )

emx_onnx_cgen/lowering/rms_normalization.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import RMSNormalizationOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import ensure_output_shape_matches_input
+from ..validation import normalize_axis
+from .common import node_dtype, shape_product, value_shape
+from .registry import register_lowering
+
+
+def _ensure_broadcastable(
+    name: str,
+    shape: tuple[int, ...],
+    normalized_shape: tuple[int, ...],
+) -> None:
+    if len(shape) != len(normalized_shape):
+        raise ShapeInferenceError(
+            f"RMSNormalization {name} rank must match normalized rank"
+        )
+    for dim, expected in zip(shape, normalized_shape):
+        if dim not in {1, expected}:
+            raise ShapeInferenceError(
+                f"RMSNormalization {name} shape {shape} must be broadcastable "
+                f"to {normalized_shape}"
+            )
+
+
+@register_lowering("RMSNormalization")
+def lower_rms_normalization(
+    graph: Graph, node: Node
+) -> RMSNormalizationOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("RMSNormalization must have 2 inputs and 1 output")
+    op_dtype = node_dtype(graph, node, *node.inputs, *node.outputs)
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "RMSNormalization 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)
+    axis = normalize_axis(int(node.attrs.get("axis", -1)), input_shape, node)
+    normalized_shape = input_shape[axis:]
+    scale_shape = value_shape(graph, node.inputs[1], node)
+    _ensure_broadcastable("scale", scale_shape, normalized_shape)
+    epsilon = float(node.attrs.get("epsilon", 1e-5))
+    stash_type = int(node.attrs.get("stash_type", 1))
+    if stash_type != 1:
+        raise UnsupportedOpError(
+            "RMSNormalization supports stash_type=1 only"
+        )
+    outer = shape_product(input_shape[:axis]) if axis > 0 else 1
+    inner = shape_product(normalized_shape)
+    return RMSNormalizationOp(
+        input0=node.inputs[0],
+        scale=node.inputs[1],
+        output=node.outputs[0],
+        shape=input_shape,
+        normalized_shape=normalized_shape,
+        scale_shape=scale_shape,
+        outer=outer,
+        inner=inner,
+        axis=axis,
+        epsilon=epsilon,
+        dtype=op_dtype,
+    )