[Bug][Relax][ONNX] ThresholdedRelu multiply-mask lowering changes nonfinite values and hardcodes FLOAT #20851

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 sources(model, feeds, expected):
    artifacts = Path("artifacts")
    artifacts.mkdir(exist_ok=True)
    identity = hashlib.sha256(model.SerializeToString()).hexdigest()
    onnx.save(model, artifacts / (identity + ".onnx"))
    np.savez(artifacts / (identity + "-inputs.npz"), **feeds)
    options = ort.SessionOptions()
    options.intra_op_num_threads = options.inter_op_num_threads = 1
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    actual = ort.InferenceSession(model.SerializeToString(), options,
                                 providers=["CPUExecutionProvider"]).run(None, feeds)
    reference = ReferenceEvaluator(model).run(None, feeds)
    for values in (actual, reference):
        for a, b in zip(values, expected, strict=True):
            np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)


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]


x = np.array([np.nan, -np.inf, -2, 0, 0.5, 1, 2, np.inf], dtype="float32")
expected = np.where(x > np.float32(1), x, np.float32(0))
for select_control in (False, True):
    if select_control:
        nodes = [h.make_node("Greater", ["x", "alpha"], ["mask"]),
                 h.make_node("Where", ["mask", "x", "zero"], ["y"])]
        constants = [nh.from_array(np.array(1, dtype="float32"), "alpha"),
                     nh.from_array(np.array(0, dtype="float32"), "zero")]
    else:
        nodes = [h.make_node("ThresholdedRelu", ["x"], ["y"], alpha=1.0)]
        constants = []
    model = make(nodes, [h.make_tensor_value_info("x", T.FLOAT, [8])],
                 [h.make_tensor_value_info("y", T.FLOAT, [8])], constants)
    sources(model, {"x": x}, [expected])
    actual = tvm_run(model, {"x": x})[0]
    print("Where control" if select_control else "ThresholdedRelu",
          "expected", expected.tolist(), "TVM", actual.tolist())
    if select_control:
        np.testing.assert_array_equal(actual, expected)
    else:
        assert np.isnan(actual[:2]).all() and (expected[:2] == 0).all()
        np.testing.assert_array_equal(actual[2:], expected[2:])

# The same converter also hardcodes FLOAT constants/masks for other legal dtypes.
for dtype, element_type in (("float16", T.FLOAT16), ("float64", T.DOUBLE)):
    data = np.array([0.5, 2.0], dtype=dtype)
    expected = np.array([0.0, 2.0], dtype=dtype)
    inputs = [h.make_tensor_value_info("x", element_type, [2])]
    outputs = [h.make_tensor_value_info("y", element_type, [2])]
    model = make([h.make_node("ThresholdedRelu", ["x"], ["y"], alpha=1.0)], inputs, outputs)
    np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, {"x": data})[0], expected)
    onnx.save(model, Path("artifacts") / (dtype + ".onnx"))
    np.savez(Path("artifacts") / (dtype + "-inputs.npz"), x=data)
    try:
        from_onnx(model, keep_params_in_input=False)
    except Exception as error:
        assert "datatype" in str(error), str(error)
        print(dtype, "conversion fails:", str(error).splitlines()[-1])
    else:
        raise AssertionError("Dtype failure did not reproduce")
    control = make([h.make_node("Greater", ["x", "alpha"], ["mask"]),
                    h.make_node("Where", ["mask", "x", "zero"], ["y"])], inputs, outputs,
                   [nh.from_array(np.array(1, dtype=dtype), "alpha"),
                    nh.from_array(np.array(0, dtype=dtype), "zero")])
    np.testing.assert_array_equal(tvm_run(control, {"x": data})[0], expected)
    onnx.save(control, Path("artifacts") / (dtype + "-control.onnx"))
    print(dtype, "typed Where control passed")