[Bug][Relax][ONNX] Resize silently drops non-unit batch scaling #20865

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)


x = np.arange(32, dtype="float32").reshape(4, 2, 2, 2)
for nc_scale in (0.3, 1.0):
    scales = np.array([nc_scale, 1, 0.7, 2.5], dtype="float32")
    shape = [1 if nc_scale == 0.3 else 4, 2, 1, 5]
    model = make([h.make_node("Resize", ["x", "roi", "scales"], ["y"],
        mode="nearest", coordinate_transformation_mode="asymmetric", nearest_mode="floor")],
        [h.make_tensor_value_info("x", T.FLOAT, [4, 2, 2, 2])],
        [h.make_tensor_value_info("y", T.FLOAT, shape)],
        [nh.from_array(np.zeros(8, dtype="float32"), "roi"), nh.from_array(scales, "scales")], opset=11)
    # ONNX asymmetric/floor indexing uses the supplied scale on every axis.
    expected = x
    for axis, (size, scale) in enumerate(zip(shape, scales)):
        indices = np.floor(np.arange(size) / scale).astype("int64")
        expected = np.take(expected, indices, axis=axis)
    save(str(nc_scale), model, {"x": x})
    reference(model, {"x": x}, [expected])
    actual = tvm_run(model, {"x": x})[0]
    print("batch scale", nc_scale, "expected shape", expected.shape, "TVM shape", actual.shape)
    if nc_scale == 0.3:
        assert actual.shape == (4, 2, 1, 5) and expected.shape == (1, 2, 1, 5)
    else:
        same(actual, expected)
        print("batch/channel scale-one control passed")