This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 59c14f6d45 [Relax][ONNX] Add GridSample ONNX frontend integration
(#18932)
59c14f6d45 is described below
commit 59c14f6d45a961a5384a7f40e52e58d2c83ea343
Author: Kryptonite <[email protected]>
AuthorDate: Thu Mar 26 03:28:26 2026 +0300
[Relax][ONNX] Add GridSample ONNX frontend integration (#18932)
### Summary
- Implements ONNX `GridSample` frontend integration for Relax, which was
previously commented out in the converter map.
- Adds `GridSample` converter class that handles ONNX→TVM grid shape
transpose (`[N, H_out, W_out, 2]` → `[N, 2, H_out, W_out]`) and
mode/padding attribute mapping.
- Reuses the existing `grid_sample` Relax op
(`relax.op.image.grid_sample`), which already exists, keeping the change
minimal and focused on the frontend layer.
- Adds tests covering all supported mode/padding_mode/align_corners
combinations.
Closes part of #18928
### Notes for Maintainers
The ONNX spec defines the default mode as `"linear"`, but onnxruntime
only accepts `"bilinear"`. I've set the converter default to
`"bilinear"` — happy to add a `"linear"` → `"bilinear"` translation if
needed for spec compliance. (**Edit: This was addressed**)
### Test Plan
```bash
python3 -m pytest -q tests/python/relax/test_frontend_onnx.py -k
'grid_sample'
```
---------
Signed-off-by: OmarAzizi <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 40 ++++++++-
tests/python/relax/test_frontend_onnx.py | 112 ++++++++++++++++++++++++
2 files changed, 151 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index edcb5e99ad..f08505951d 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -3896,6 +3896,44 @@ class AllClassNMS(OnnxOpConverter):
return nms_out
+class GridSample(OnnxOpConverter):
+ """Converts an onnx GridSample node into an equivalent Relax expression."""
+
+ @classmethod
+ def _impl_v16(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ grid = inputs[1]
+
+ method = attr.get("mode", b"bilinear")
+ if isinstance(method, bytes):
+ method = method.decode("ascii")
+
+ # Translate ONNX mode names to TVM method names
+ if method == "linear":
+ method = "bilinear"
+ elif method == "cubic":
+ method = "bicubic"
+
+ padding_mode = attr.get("padding_mode", b"zeros")
+ if isinstance(padding_mode, bytes):
+ padding_mode = padding_mode.decode("ascii")
+
+ align_corners = bool(attr.get("align_corners", 0))
+
+ # ONNX grid shape: [N, H_out, W_out, 2]
+ # TVM grid shape: [N, 2, H_out, W_out]
+ grid = relax.op.permute_dims(grid, [0, 3, 1, 2])
+
+ return relax.op.image.grid_sample(
+ data,
+ grid,
+ method=method,
+ layout="NCHW",
+ padding_mode=padding_mode,
+ align_corners=align_corners,
+ )
+
+
def _get_convert_map():
return {
# defs/experimental
@@ -4048,7 +4086,7 @@ def _get_convert_map():
# "RoiAlign": RoiAlign,
"NonMaxSuppression": NonMaxSuppression,
"AllClassNMS": AllClassNMS,
- # "GridSample": GridSample,
+ "GridSample": GridSample,
"Upsample": Upsample,
# others
"DepthToSpace": DepthToSpace,
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 7ea80c1bbe..8740720205 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3979,6 +3979,118 @@ def test_nms_score_threshold():
tvm_selected[:min_rows], ort_selected[:min_rows], rtol=1e-5,
atol=1e-5
)
[email protected]("mode", ["bilinear", "nearest", "bicubic"])
[email protected]("padding_mode", ["zeros", "border", "reflection"])
[email protected]("align_corners", [0, 1])
+def test_grid_sample(mode, padding_mode, align_corners):
+ # Only testing 2D (NCHW) as that's what TVM currently supports
+ x_shape = [1, 3, 4, 4]
+ grid_shape = [1, 2, 2, 2]
+ out_shape = [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2]]
+
+ node = helper.make_node(
+ "GridSample",
+ inputs=["X", "grid"],
+ outputs=["Y"],
+ mode=mode,
+ padding_mode=padding_mode,
+ align_corners=align_corners,
+ )
+
+ graph = helper.make_graph(
+ [node],
+ "grid_sample_test",
+ inputs=[
+ helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
+ helper.make_tensor_value_info("grid", TensorProto.FLOAT,
grid_shape),
+ ],
+ outputs=[
+ helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape),
+ ],
+ )
+
+ # Grid values must be in [-1, 1]: -1 is far left/top, 1 is far right/bottom
+ grid_data = np.random.uniform(-1, 1, grid_shape).astype("float32")
+ # Use controlled X input to avoid extreme values affecting nearest mode
boundaries
+ x_data = np.random.uniform(-1, 1, x_shape).astype("float32")
+
+ model = helper.make_model(graph, producer_name="grid_sample_test")
+ check_correctness(
+ model,
+ inputs={"grid": grid_data, "X": x_data},
+ opset=16,
+ )
+
+def test_grid_sample_linear_mode_translation():
+ """Test that ONNX mode='linear' is correctly translated to 'bilinear'.
+
+ The ONNX spec defines 'linear' as a valid mode for GridSample, but
+ onnxruntime rejects it in practice. Real ONNX models exported from
+ frameworks like PyTorch may still use 'linear'. We verify the translation
+ by inspecting the Relax IR directly rather than running check_correctness.
+ """
+ x_shape = [1, 3, 4, 4]
+ grid_shape = [1, 2, 2, 2]
+
+ node = helper.make_node(
+ "GridSample",
+ inputs=["X", "grid"],
+ outputs=["Y"],
+ mode="linear",
+ )
+
+ graph = helper.make_graph(
+ [node],
+ "grid_sample_linear_test",
+ inputs=[
+ helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
+ helper.make_tensor_value_info("grid", TensorProto.FLOAT,
grid_shape),
+ ],
+ outputs=[
+ helper.make_tensor_value_info("Y", TensorProto.FLOAT, [x_shape[0],
x_shape[1], grid_shape[1], grid_shape[2]]),
+ ],
+ )
+
+ model = helper.make_model(graph, producer_name="grid_sample_linear_test")
+ tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)
+ # Verify 'linear' was translated to 'bilinear' in the Relax IR
+ assert 'method="bilinear"' in str(tvm_model)
+
+
+def test_grid_sample_cubic_mode_translation():
+ """Test that ONNX mode='cubic' is correctly translated to 'bicubic'.
+
+ The ONNX spec defines 'cubic' as a valid mode for GridSample, but
+ TVM uses 'bicubic'. We verify the translation by inspecting the
+ Relax IR directly rather than running check_correctness.
+ """
+ x_shape = [1, 3, 4, 4]
+ grid_shape = [1, 2, 2, 2]
+
+ node = helper.make_node(
+ "GridSample",
+ inputs=["X", "grid"],
+ outputs=["Y"],
+ mode="cubic",
+ )
+
+ graph = helper.make_graph(
+ [node],
+ "grid_sample_cubic_test",
+ inputs=[
+ helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
+ helper.make_tensor_value_info("grid", TensorProto.FLOAT,
grid_shape),
+ ],
+ outputs=[
+ helper.make_tensor_value_info("Y", TensorProto.FLOAT, [x_shape[0],
x_shape[1], grid_shape[1], grid_shape[2]]),
+ ],
+ )
+
+ model = helper.make_model(graph, producer_name="grid_sample_cubic_test")
+ tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)
+ # Verify 'cubic' was translated to 'bicubic' in the Relax IR
+ assert 'method="bicubic"' in str(tvm_model)
if __name__ == "__main__":
tvm.testing.main()
+
