[Bug][Relax][ONNX] ReduceLogSumExp produces NaN for a positive-infinity input #20852

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]


for opset in (13, 18):
    attrs = {"keepdims": 1}
    if opset == 13:
        attrs["axes"] = [0]
    model = make([h.make_node("ReduceLogSumExp", ["x"] if opset == 13 else ["x", "axes"],
                              ["y"], **attrs)],
                 [h.make_tensor_value_info("x", T.FLOAT, [3])],
                 [h.make_tensor_value_info("y", T.FLOAT, [1])],
                 [] if opset == 13 else [nh.from_array(np.array([0], dtype="int64"), "axes")],
                 opset=opset)
    control = make([
        h.make_node("Cast", ["x"], ["dx"], to=T.DOUBLE),
        h.make_node("Exp", ["dx"], ["ex"]),
        h.make_node("ReduceSum", ["ex", "axes"], ["sum"], keepdims=1),
        h.make_node("Log", ["sum"], ["log"]),
        h.make_node("Cast", ["log"], ["y"], to=T.FLOAT)],
        [h.make_tensor_value_info("x", T.FLOAT, [3])],
        [h.make_tensor_value_info("y", T.FLOAT, [1])],
        [nh.from_array(np.array([0], dtype="int64"), "axes")], opset=opset)
    for values in ([float("inf"), 0.0, 1.0], [-1.0, 0.0, 1.0]):
        x = np.array(values, dtype="float32")
        expected = np.logaddexp.reduce(x).reshape(1)
        sources(model, {"x": x}, [expected])
        actual = tvm_run(model, {"x": x})[0]
        print("opset", opset, "input", x.tolist(), "expected", expected.tolist(),
              "TVM", actual.tolist())
        if np.isposinf(x).any():
            assert np.isposinf(expected).all() and np.isnan(actual).all()
        else:
            np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6)
        sources(control, {"x": x}, [expected])
        expanded = tvm_run(control, {"x": x})[0]
        np.testing.assert_allclose(expanded, expected, rtol=1e-5, atol=1e-6)
        print("expanded ONNX function control", expanded.tolist())