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

emx_onnx_cgen/lowering/expand.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import ExpandOp
+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_shape_values(graph: Graph, name: str, node: Node) -> 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(
+            f"{node.op_type} shape input must be int64 or int32"
+        )
+    if len(initializer.type.shape) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} shape input must be a 1D tensor"
+        )
+    values = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if values.size == 0:
+        raise ShapeInferenceError(
+            f"{node.op_type} shape input cannot be empty"
+        )
+    return [int(value) for value in values]
+
+
+def _validate_shape_input(graph: Graph, name: str, node: Node) -> None:
+    dtype = value_dtype(graph, name, node)
+    if dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            f"{node.op_type} shape input must be int64 or int32"
+        )
+    shape = value_shape(graph, name, node)
+    if len(shape) != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} shape input must be a 1D tensor"
+        )
+    if shape[0] <= 0:
+        raise ShapeInferenceError(
+            f"{node.op_type} shape input cannot be empty"
+        )
+
+
+def _validate_static_dims(shape: tuple[int, ...], node: Node) -> None:
+    if any(dim < 0 for dim in shape):
+        raise ShapeInferenceError(
+            f"{node.op_type} does not support dynamic dims"
+        )
+
+
+def _broadcast_shape(
+    input_shape: tuple[int, ...], shape_values: list[int], node: Node
+) -> tuple[int, ...]:
+    _validate_static_dims(input_shape, node)
+    for dim in shape_values:
+        if dim < 0:
+            raise ShapeInferenceError(
+                f"{node.op_type} does not support dynamic dims"
+            )
+    output_rank = max(len(input_shape), len(shape_values))
+    input_padded = (1,) * (output_rank - len(input_shape)) + input_shape
+    shape_padded = (1,) * (output_rank - len(shape_values)) + tuple(shape_values)
+    result: list[int] = []
+    for input_dim, shape_dim in zip(input_padded, shape_padded):
+        if input_dim == 1:
+            result.append(shape_dim)
+        elif shape_dim == 1:
+            result.append(input_dim)
+        elif input_dim == shape_dim:
+            result.append(input_dim)
+        else:
+            raise ShapeInferenceError(
+                f"{node.op_type} input shape {input_shape} is not "
+                f"broadcastable to {shape_values}"
+            )
+    return tuple(result)
+
+
+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("Expand")
+def lower_expand(graph: Graph, node: Node) -> ExpandOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Expand must have 2 inputs and 1 output")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[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(
+            f"{node.op_type} expects matching input/output dtypes, "
+            f"got {input_dtype} and {output_dtype}"
+        )
+    shape_values = _read_shape_values(graph, node.inputs[1], node)
+    if shape_values is not None:
+        expected_output_shape = _broadcast_shape(input_shape, shape_values, node)
+        _validate_static_dims(expected_output_shape, node)
+        if output_shape and output_shape != expected_output_shape:
+            raise ShapeInferenceError(
+                f"{node.op_type} output shape must be {expected_output_shape}, "
+                f"got {output_shape}"
+            )
+    else:
+        _validate_shape_input(graph, node.inputs[1], node)
+        if not output_shape:
+            raise ShapeInferenceError(
+                f"{node.op_type} output shape must be specified"
+            )
+        expected_output_shape = _broadcast_shape(
+            input_shape, list(output_shape), node
+        )
+        if expected_output_shape != output_shape:
+            raise ShapeInferenceError(
+                f"{node.op_type} output shape must be {expected_output_shape}, "
+                f"got {output_shape}"
+            )
+    input_shape_padded = (
+        (1,) * (len(expected_output_shape) - len(input_shape)) + input_shape
+    )
+    input_strides = _compute_strides(input_shape_padded)
+    return ExpandOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=expected_output_shape,
+        input_shape_padded=input_shape_padded,
+        input_strides=input_strides,
+        dtype=input_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/eye_like.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import EyeLikeOp
+from ..dtypes import scalar_type_from_onnx
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+@register_lowering("EyeLike")
+def lower_eye_like(graph: Graph, node: Node) -> EyeLikeOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("EyeLike must have 1 input and 1 output")
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if input_shape != output_shape:
+        raise ShapeInferenceError("EyeLike input and output shapes must match")
+    if len(output_shape) < 2:
+        raise UnsupportedOpError("EyeLike expects input rank >= 2")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    dtype_attr = node.attrs.get("dtype")
+    if dtype_attr is not None:
+        target_dtype = scalar_type_from_onnx(int(dtype_attr))
+        if target_dtype is None:
+            raise UnsupportedOpError(
+                f"EyeLike dtype {dtype_attr} is not supported"
+            )
+        if output_dtype != target_dtype:
+            raise UnsupportedOpError(
+                "EyeLike output dtype must match dtype attribute, "
+                f"got {output_dtype.onnx_name} and {target_dtype.onnx_name}"
+            )
+    k = int(node.attrs.get("k", 0))
+    return EyeLikeOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        output_shape=output_shape,
+        k=k,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/flatten.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from ..codegen.c_emitter import ReshapeOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import shape_product, value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _normalize_axis(axis: int, rank: int) -> int:
+    if axis < 0:
+        axis += rank
+    if axis < 0 or axis > rank:
+        raise UnsupportedOpError("Flatten axis must be within input rank")
+    return axis
+
+
+def _flatten_output_shape(
+    input_shape: tuple[int, ...], axis: int
+) -> tuple[int, int]:
+    rank = len(input_shape)
+    axis = _normalize_axis(axis, rank)
+    if rank == 0:
+        return (1, 1)
+    for dim in input_shape:
+        if dim < 0:
+            raise ShapeInferenceError("Dynamic dims are not supported")
+    first = shape_product(input_shape[:axis]) if axis else 1
+    second = shape_product(input_shape[axis:]) if axis < rank else 1
+    return (first, second)
+
+
+@register_lowering("Flatten")
+def lower_flatten(graph: Graph, node: Node) -> ReshapeOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Flatten must have 1 input 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(
+            "Flatten expects matching input/output dtypes, "
+            f"got {input_dtype} and {output_dtype}"
+        )
+    axis = int(node.attrs.get("axis", 1))
+    output_shape = _flatten_output_shape(input_shape, axis)
+    expected_shape = value_shape(graph, node.outputs[0], node)
+    if expected_shape and output_shape != expected_shape:
+        raise ShapeInferenceError(
+            "Flatten output shape must be "
+            f"{output_shape}, got {expected_shape}"
+        )
+    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/gather.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import GatherOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import normalize_axis
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@register_lowering("Gather")
+def lower_gather(graph: Graph, node: Node) -> GatherOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("Gather must have 2 inputs and 1 output")
+    data_name, indices_name = node.inputs
+    data_shape = _value_shape(graph, data_name, node)
+    indices_shape = _value_shape(graph, indices_name, node)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    axis = normalize_axis(int(node.attrs.get("axis", 0)), data_shape, node)
+    expected_output_shape = (
+        data_shape[:axis] + indices_shape + data_shape[axis + 1 :]
+    )
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "Gather output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    op_dtype = _value_dtype(graph, data_name, node)
+    indices_dtype = _value_dtype(graph, indices_name, node)
+    if indices_dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            "Gather indices must be int32 or int64, "
+            f"got {indices_dtype.onnx_name}"
+        )
+    return GatherOp(
+        data=data_name,
+        indices=indices_name,
+        output=node.outputs[0],
+        axis=axis,
+        data_shape=data_shape,
+        indices_shape=indices_shape,
+        output_shape=output_shape,
+        dtype=op_dtype,
+        indices_dtype=indices_dtype,
+    )

emx_onnx_cgen/lowering/gather_elements.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import GatherElementsOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import normalize_axis
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@register_lowering("GatherElements")
+def lower_gather_elements(graph: Graph, node: Node) -> GatherElementsOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("GatherElements must have 2 inputs and 1 output")
+    data_name, indices_name = node.inputs
+    data_shape = _value_shape(graph, data_name, node)
+    indices_shape = _value_shape(graph, indices_name, node)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    if len(data_shape) != len(indices_shape):
+        raise ShapeInferenceError(
+            "GatherElements inputs must have matching ranks, "
+            f"got {data_shape} and {indices_shape}"
+        )
+    if output_shape != indices_shape:
+        raise ShapeInferenceError(
+            "GatherElements output shape must match indices shape, "
+            f"got {output_shape} and {indices_shape}"
+        )
+    axis = normalize_axis(int(node.attrs.get("axis", 0)), data_shape, node)
+    for dim_index, (data_dim, index_dim) in enumerate(
+        zip(data_shape, indices_shape)
+    ):
+        if dim_index == axis:
+            continue
+        if data_dim != index_dim:
+            raise ShapeInferenceError(
+                "GatherElements inputs must match on non-axis dimensions, "
+                f"got {data_shape} and {indices_shape}"
+            )
+    op_dtype = _value_dtype(graph, data_name, node)
+    indices_dtype = _value_dtype(graph, indices_name, node)
+    if indices_dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            "GatherElements indices must be int32 or int64, "
+            f"got {indices_dtype.onnx_name}"
+        )
+    return GatherElementsOp(
+        data=data_name,
+        indices=indices_name,
+        output=node.outputs[0],
+        axis=axis,
+        data_shape=data_shape,
+        indices_shape=indices_shape,
+        output_shape=output_shape,
+        dtype=op_dtype,
+        indices_dtype=indices_dtype,
+    )

emx_onnx_cgen/lowering/gemm.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import GemmOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import node_dtype as _node_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class GemmSpec:
+    m: int
+    n: int
+    k: int
+    alpha: float | int
+    beta: float | int
+    trans_a: bool
+    trans_b: bool
+    c_shape: tuple[int, ...] | None
+
+
+def resolve_gemm_spec(graph: Graph, node: Node, dtype: ScalarType) -> GemmSpec:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError("Gemm must have 2 or 3 inputs and 1 output")
+    alpha, beta, trans_a, trans_b = _resolve_gemm_attrs(node, dtype)
+    input0_shape = _value_shape(graph, node.inputs[0], node)
+    input1_shape = _value_shape(graph, node.inputs[1], node)
+    if len(input0_shape) != 2 or len(input1_shape) != 2:
+        raise UnsupportedOpError(
+            "Gemm supports 2D inputs only, "
+            f"got {input0_shape} x {input1_shape}"
+        )
+    if trans_a:
+        m, k_left = input0_shape[1], input0_shape[0]
+    else:
+        m, k_left = input0_shape
+    if trans_b:
+        n, k_right = input1_shape[0], input1_shape[1]
+    else:
+        k_right, n = input1_shape
+    if k_left != k_right:
+        raise ShapeInferenceError(
+            f"Gemm inner dimensions must match, got {k_left} and {k_right}"
+        )
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    if output_shape != (m, n):
+        raise ShapeInferenceError(
+            f"Gemm output shape must be {(m, n)}, got {output_shape}"
+        )
+    c_shape = None
+    if len(node.inputs) == 3:
+        bias_shape = _value_shape(graph, node.inputs[2], node)
+        c_shape = validate_gemm_bias_shape((m, n), bias_shape, node)
+    return GemmSpec(
+        m=m,
+        n=n,
+        k=k_left,
+        alpha=alpha,
+        beta=beta,
+        trans_a=trans_a,
+        trans_b=trans_b,
+        c_shape=c_shape,
+    )
+
+
+def _resolve_gemm_attrs(
+    node: Node, dtype: ScalarType
+) -> tuple[float | int, float | int, bool, bool]:
+    alpha = float(node.attrs.get("alpha", 1.0))
+    beta = float(node.attrs.get("beta", 1.0))
+    trans_a = int(node.attrs.get("transA", 0))
+    trans_b = int(node.attrs.get("transB", 0))
+    if trans_a not in {0, 1} or trans_b not in {0, 1}:
+        raise UnsupportedOpError(
+            "Gemm only supports transA/transB values of 0 or 1"
+        )
+    if dtype == ScalarType.BOOL:
+        raise UnsupportedOpError("Gemm supports numeric inputs only")
+    if not dtype.is_float:
+        alpha_int = int(alpha)
+        beta_int = int(beta)
+        if alpha != alpha_int or beta != beta_int:
+            raise UnsupportedOpError(
+                "Gemm alpha and beta must be integers for non-float inputs"
+            )
+        alpha = alpha_int
+        beta = beta_int
+    return alpha, beta, bool(trans_a), bool(trans_b)
+
+
+def validate_gemm_bias_shape(
+    output_shape: tuple[int, int], bias_shape: tuple[int, ...], node: Node
+) -> tuple[int, ...]:
+    if len(bias_shape) == 0:
+        return bias_shape
+    if len(bias_shape) == 1:
+        if bias_shape[0] not in {1, output_shape[1]}:
+            raise ShapeInferenceError(
+                "Gemm bias input must be broadcastable to output shape, "
+                f"got {bias_shape} vs {output_shape}"
+            )
+        return bias_shape
+    if len(bias_shape) == 2:
+        m, n = output_shape
+        if bias_shape[0] not in {1, m} or bias_shape[1] not in {1, n}:
+            raise ShapeInferenceError(
+                "Gemm bias input must be broadcastable to output shape, "
+                f"got {bias_shape} vs {output_shape}"
+            )
+        return bias_shape
+    raise ShapeInferenceError(
+        f"Gemm bias input must be rank 1 or 2, got {bias_shape}"
+    )
+
+
+@register_lowering("Gemm")
+def lower_gemm(graph: Graph, node: Node) -> GemmOp:
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    spec = resolve_gemm_spec(graph, node, op_dtype)
+    return GemmOp(
+        input_a=node.inputs[0],
+        input_b=node.inputs[1],
+        input_c=node.inputs[2] if len(node.inputs) == 3 else None,
+        output=node.outputs[0],
+        m=spec.m,
+        n=spec.n,
+        k=spec.k,
+        trans_a=spec.trans_a,
+        trans_b=spec.trans_b,
+        alpha=spec.alpha,
+        beta=spec.beta,
+        c_shape=spec.c_shape,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/grid_sample.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import GridSampleOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
+from .registry import register_lowering
+
+_SUPPORTED_MODES = {"linear", "nearest", "cubic"}
+_SUPPORTED_PADDING_MODES = {"zeros", "border", "reflection"}
+
+
+@dataclass(frozen=True)
+class _GridSampleShapes:
+    input_shape: tuple[int, ...]
+    grid_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    spatial_rank: int
+
+
+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 _resolve_shapes(graph: Graph, node: Node) -> _GridSampleShapes:
+    input_shape = value_shape(graph, node.inputs[0], node)
+    grid_shape = value_shape(graph, node.inputs[1], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if len(input_shape) < 3:
+        raise ShapeInferenceError(
+            "GridSample expects input rank of at least 3"
+        )
+    spatial_rank = len(input_shape) - 2
+    if any(dim < 0 for dim in (*input_shape, *grid_shape, *output_shape)):
+        raise ShapeInferenceError(
+            "GridSample requires static, non-negative shapes"
+        )
+    return _GridSampleShapes(
+        input_shape=input_shape,
+        grid_shape=grid_shape,
+        output_shape=output_shape,
+        spatial_rank=spatial_rank,
+    )
+
+
+def _validate_shapes(shapes: _GridSampleShapes) -> None:
+    input_shape = shapes.input_shape
+    grid_shape = shapes.grid_shape
+    output_shape = shapes.output_shape
+    spatial_rank = shapes.spatial_rank
+    if len(grid_shape) != spatial_rank + 2:
+        raise ShapeInferenceError(
+            "GridSample expects grid rank to match input spatial rank"
+        )
+    if len(output_shape) != spatial_rank + 2:
+        raise ShapeInferenceError(
+            "GridSample expects output rank to match input spatial rank"
+        )
+    if grid_shape[0] != input_shape[0]:
+        raise ShapeInferenceError("GridSample expects matching batch dimension")
+    if grid_shape[-1] != spatial_rank:
+        raise ShapeInferenceError(
+            "GridSample expects grid last dimension to match spatial rank"
+        )
+    expected_output = (
+        input_shape[0],
+        input_shape[1],
+        *grid_shape[1:-1],
+    )
+    if output_shape != expected_output:
+        raise ShapeInferenceError(
+            "GridSample output shape must be "
+            f"{expected_output}, got {output_shape}"
+        )
+
+
+def _validate_dtypes(
+    graph: Graph, node: Node
+) -> tuple[ScalarType, ScalarType]:
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    grid_dtype = value_dtype(graph, node.inputs[1], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "GridSample expects matching input/output dtypes, got "
+            f"{input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if not input_dtype.is_float:
+        raise UnsupportedOpError(
+            "GridSample currently supports floating-point inputs only"
+        )
+    if not grid_dtype.is_float:
+        raise UnsupportedOpError("GridSample expects floating-point grid")
+    return input_dtype, grid_dtype
+
+
+@register_lowering("GridSample")
+def lower_grid_sample(graph: Graph, node: Node) -> GridSampleOp:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "GridSample expects 2 inputs (X, grid) and 1 output"
+        )
+    shapes = _resolve_shapes(graph, node)
+    _validate_shapes(shapes)
+    mode = _decode_attr(node.attrs.get("mode"), "linear")
+    padding_mode = _decode_attr(node.attrs.get("padding_mode"), "zeros")
+    align_corners = int(node.attrs.get("align_corners", 0))
+    if mode not in _SUPPORTED_MODES:
+        raise UnsupportedOpError(
+            f"GridSample mode {mode!r} is not supported"
+        )
+    if padding_mode not in _SUPPORTED_PADDING_MODES:
+        raise UnsupportedOpError(
+            "GridSample padding_mode "
+            f"{padding_mode!r} is not supported"
+        )
+    if align_corners not in {0, 1}:
+        raise UnsupportedOpError("GridSample align_corners must be 0 or 1")
+    input_dtype, grid_dtype = _validate_dtypes(graph, node)
+    if shapes.spatial_rank > 3:
+        raise UnsupportedOpError(
+            "GridSample supports up to 3 spatial dimensions"
+        )
+    return GridSampleOp(
+        input0=node.inputs[0],
+        grid=node.inputs[1],
+        output=node.outputs[0],
+        input_shape=shapes.input_shape,
+        grid_shape=shapes.grid_shape,
+        output_shape=shapes.output_shape,
+        spatial_rank=shapes.spatial_rank,
+        input_spatial=shapes.input_shape[2:],
+        output_spatial=shapes.output_shape[2:],
+        mode=mode,
+        padding_mode=padding_mode,
+        align_corners=bool(align_corners),
+        dtype=input_dtype,
+        grid_dtype=grid_dtype,
+    )