[Bug][Relax][ONNX] Constant GridSample with DOUBLE coordinates changes a FLOAT result to DOUBLE

Start by running the attached repro.py to confirm the double_grid case fails with the dtype mismatch while float_grid passes. The GridSample converter is in python/tvm/relax/frontend/onnx/onnx_frontend.py, and the constant conversion/folding path that produces the sample is where the coordinate dtype leaks into the result. Done when the double-grid sample keeps the FLOAT image dtype, the Add runs, and the output is float32 with value 2.5.

The image and coordinate types are allowed to differ; the output type should still follow the image.

A constant FLOAT image and DOUBLE sampling grid are legal GridSample inputs. Conversion/folding produces a DOUBLE sample; consuming it with a FLOAT Add fails the same-dtype check. Replacing only the coordinate tensor's dtype with FLOAT makes the graph run and return 2.5.

A runnable case

Save the code below as repro.py. With an LLVM-enabled TVM installation and the dependencies listed at the end, run OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 python repro.py.

from pathlib import Path
import hashlib
import numpy as np
import onnx
import onnxruntime as ort
import tvm
from onnx import TensorProto as T, helper as h, numpy_helper as nh
from onnx.reference import ReferenceEvaluator
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx


def make(nodes, inputs, outputs, initializers=(), opset=18):
    model = h.make_model(h.make_graph(nodes, "repro", inputs, outputs, list(initializers)),
                         opset_imports=[h.make_opsetid("", opset)], ir_version=10)
    onnx.checker.check_model(model, full_check=True)
    return model


def tvm_run(model, feeds):
    mod = from_onnx(model, keep_params_in_input=False)
    names = [value.name for value in model.graph.input]
    executable = relax.build(mod, target="llvm", relax_pipeline="default")
    result = relax.VirtualMachine(executable, tvm.cpu())["main"](
        *[tvm.runtime.tensor(feeds[name]) for name in names])
    return [result.numpy()] if hasattr(result, "numpy") else [v.numpy() for v in result]


def save(label, model, feeds):
    folder = Path("artifacts") / label
    folder.mkdir(parents=True, exist_ok=True)
    onnx.checker.check_model(model, full_check=True)
    onnx.save(model, folder / "model.onnx")
    np.savez(folder / "inputs.npz", **feeds)
    print(label, "model_sha256", hashlib.sha256((folder / "model.onnx").read_bytes()).hexdigest())
    print(label, "inputs_sha256", hashlib.sha256((folder / "inputs.npz").read_bytes()).hexdigest())


def reference(model, feeds, expected, check_ort=True):
    outputs = ReferenceEvaluator(model).run(None, feeds)
    for a, b in zip(outputs, expected, strict=True):
        assert a.shape == b.shape and a.dtype == b.dtype, (a.shape, a.dtype, b.shape, b.dtype)
        np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)
    if check_ort:
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = opts.inter_op_num_threads = 1
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
        outputs = ort.InferenceSession(model.SerializeToString(), opts,
            providers=["CPUExecutionProvider"]).run(None, feeds)
        for a, b in zip(outputs, expected, strict=True):
            assert a.shape == b.shape and a.dtype == b.dtype
            np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)


def error_line(error):
    return type(error).__name__ + ": " + str(error).splitlines()[-1]


def same(actual, expected):
    assert actual.shape == expected.shape and actual.dtype == expected.dtype
    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6)


for label, grid_dtype in (("double_grid", "float64"), ("float_grid", "float32")):
    model = make([h.make_node("GridSample", ["x", "grid"], ["sample"],
                    mode="cubic", padding_mode="reflection", align_corners=1),
                  h.make_node("Add", ["sample", "offset"], ["y"])],
        [h.make_tensor_value_info("offset", T.FLOAT, [1, 1, 1, 1])],
        [h.make_tensor_value_info("y", T.FLOAT, [1, 1, 1, 1])],
        [nh.from_array(np.arange(1, 5, dtype="float32").reshape(1, 1, 2, 2), "x"),
         nh.from_array(np.zeros((1, 1, 1, 2), dtype=grid_dtype), "grid")], opset=20)
    feeds = {"offset": np.zeros((1, 1, 1, 1), dtype="float32")}
    save(label, model, feeds)
    # At the center, bicubic interpolation of this affine 2x2 field is 2.5.
    expected = np.full((1, 1, 1, 1), 2.5, dtype="float32")
    try:
        actual = tvm_run(model, feeds)[0]
    except Exception as error:
        assert label == "double_grid" and ("float64" in str(error) or "float32" in str(error)), str(error)
        print(label, error_line(error))
    else:
        assert label == "float_grid", "Mixed-coordinate constant-fold failure did not reproduce"
        same(actual, expected)
        print("FLOAT coordinates control passed", actual.tolist())

What I see

double_grid TypeError: Binary operators must have the same datatype for both operands.  However, ir.Op(loc=UnknownLoc(), ty=I.MissingType(), index=48, name="relax.add", doc="", args_info=(ir.ArgumentInfo(name="x1", doc="The first input tensor."), ir.ArgumentInfo(name="x2", doc="The second input tensor.")), attrs_type_key="", ty_args_info=(), var_args_info=None, var_ty_args_info=ir.ArgumentInfo(name="out_type", doc="Optional output tensor type carrying the virtual device.")) uses datatype T.float64 on the LHS (Type of R.Tensor((1, 1, 1, 1), dtype="float64")), and datatype T.float32 on the RHS (Type of R.Tensor((1, 1, 1, 1), dtype="float32")).
FLOAT coordinates control passed [[[[2.5]]]]

The expected/actual distinction is: The constant sample retains FLOAT and Add returns [[[[2.5]]]]. TVM instead reports FLOAT64 on the Add left side versus FLOAT32 on the right. FLOAT coordinates pass.

Why I think this is the relevant boundary

GridSample-20 uses T1 for the image/output and a separate floating T2 for the grid. This witness has FLOAT T1 and DOUBLE T2. At the image center, bicubic interpolation of the affine 2x2 field [[1,2],[3,4]] is 2.5; the FLOAT-grid control verifies the same geometry in TVM. No ORT or ONNX ReferenceEvaluator execution claim is made for this mixed-grid witness.

The GridSample converter accepts separate image/grid types. In the constant conversion/folding path the coordinate dtype escapes into the sample result, violating the declared output type. The failure precedes execution of the Add. A correction should preserve the image/output dtype during constant evaluation; it is not a blanket request to reject DOUBLE grids.

Scope of this example

#18932 introduced Relax GridSample frontend support; generic support does not establish correctness of this mixed-type constant path. Ordinary DOUBLE-image/DOUBLE-grid support is not the issue reported here. Targeted mixed-type searches found no exact covering report.

This self-contained script is a newly constructed minimal diagnostic from the same graph-testing investigation, not a byte-identical replay of an archived model. It needs only TVM, ONNX, NumPy and ONNX Runtime, downloads no files, and writes its generated models and inputs to artifacts/. The assertions check both the observed failure/wrong result and the normal controls; exit code zero means that diagnostic reproduced. No global novelty or current-main execution claim is made.

Reproduction setup

Linux x86_64 CPU/LLVM; Python 3.12.3; TVM 0.28.dev0; LLVM 18.1.3; ONNX 1.22.0; NumPy 2.5.4; ONNX Runtime 1.29.0 (CPUExecutionProvider).

The installed TVM binary's exact build source revision is UNKNOWN. Source links are a separate static inspection of main at 70ec7d5bf20326ef301ee3f06eb6aed51760331d; that revision was not built or executed for this report.

Loaded binary SHA256 identities:

• libtvm_ffi.so: 12072d0ccd9d5cb2c6dac80d42f4217ab50dbf795db03197000614a7a770a816
• libtvm_compiler.so: 574a573fe994e5eda7352ddd8d685e64e45c7934bdeeaab57a5a272a77792ab1
• libtvm_runtime.so: 4175916fc0ce6266c912629d2255094b94087425c6ab53ffd5e004d4c292c828
• libtvm_runtime_extra.so: 79c93351fca12dd7851964890cea7da65dd2b2e0e64e5c28bfed56c878b70817

The code and inputs are small enough to recreate without an attachment. For an exact file check, repro.py has SHA256 baa510dff58f8bff7f040456b0ce4f4dd76ab59ca6d0e4c6f099f4b213d2ea84. Generated model/input checksums from this run:

• artifacts/double_grid/inputs.npz: c0af769bd3effd1d5a3f2b72ab1dfb1837fed89911a46bd8d25c604620f9b9cc
• artifacts/double_grid/model.onnx: 6910a07f6d8af0973f068f1b25525b5d2bdcac45bbd3f1f69e2f640d587acf20
• artifacts/float_grid/inputs.npz: c0af769bd3effd1d5a3f2b72ab1dfb1837fed89911a46bd8d25c604620f9b9cc
• artifacts/float_grid/model.onnx: f82a6cd6d0c3b85cc067699a3d760c6864b6d0264a6f3f1e9a37aa477ec5b868

I can follow up with a focused regression test and adjust the example if another supported lowering path is preferred.