[Relax][Frontend][ONNX] Lower ThresholdedRelu with where in the input dtype

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -2048,7 +2048,12 @@ class ThresholdedRelu(OnnxOpConverter):
     def _impl_v1(cls, bb, inputs, attr, params):
         x = inputs[0]
         alpha = attr.get("alpha", 1.0)
-        return relax.op.greater(x, relax.const(alpha)).astype("float32") * x
+        dtype = x.ty.dtype
+        # Select rather than multiply by a mask: nan * 0 and -inf * 0 are nan, and a
+        # float32 mask would not match float16/float64 inputs.
+        return relax.op.where(
+            relax.op.greater(x, relax.const(alpha, dtype)), x, relax.const(0, dtype)
+        )
 
 
 class LeakyRelu(OnnxOpConverter):

tests/python/relax/test_frontend_onnx.py

@@ -4793,8 +4793,7 @@ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float3
             R.func_attr({"num_input": 1})
             with R.dataflow():
                 lv: R.Tensor((2, 3), dtype="bool") = R.greater(x, R.const(1.0, "float32"))
-                lv1: R.Tensor((2, 3), dtype="float32") = R.astype(lv, dtype="float32")
-                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(lv1, x)
+                gv: R.Tensor((2, 3), dtype="float32") = R.where(lv, x, R.const(0.0, "float32"))
                 R.output(gv)
             return gv
 
@@ -4812,14 +4811,47 @@ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float3
                 lv: R.Tensor((2, 3), dtype="bool") = R.greater(
                     x, R.const(-0.0099999997764825821, "float32")
                 )
-                lv1: R.Tensor((2, 3), dtype="float32") = R.astype(lv, dtype="float32")
-                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(lv1, x)
+                gv: R.Tensor((2, 3), dtype="float32") = R.where(lv, x, R.const(0.0, "float32"))
                 R.output(gv)
             return gv
 
     tvm_tirx.ir.assert_structural_equal(tvm_model, ExpectedCustom)
 
 
+@pytest.mark.parametrize(
+    "onnx_dtype, np_dtype",
+    [
+        (TensorProto.FLOAT, np.float32),
+        (TensorProto.FLOAT16, np.float16),
+        (TensorProto.DOUBLE, np.float64),
+    ],
+)
+def test_thresholded_relu_nonfinite_and_dtypes(onnx_dtype, np_dtype):
+    # nan and -inf are not above alpha, so they map to 0 (a mask multiply gave nan).
+    xs = np.array([np.nan, -np.inf, -2.0, 0.0, 0.5, 1.0, 2.0, np.inf], dtype=np_dtype)
+    expected = np.where(xs > 1.0, xs, 0).astype(np_dtype)
+    graph = helper.make_graph(
+        [helper.make_node("ThresholdedRelu", ["x"], ["y"], alpha=1.0)],
+        "thresholded_relu_nonfinite",
+        inputs=[helper.make_tensor_value_info("x", onnx_dtype, [len(xs)])],
+        outputs=[helper.make_tensor_value_info("y", onnx_dtype, [len(xs)])],
+    )
+    model = helper.make_model(graph, producer_name="thresholded_relu_nonfinite")
+
+    tvm_model = from_onnx(model, opset=10, keep_params_in_input=True)
+    tvm_model = relax.transform.LegalizeOps()(tvm_model)
+    tvm_model, _ = relax.frontend.detach_params(tvm_model)
+    with tvm_tirx.ir.transform.PassContext(opt_level=3):
+        ex = tvm.compile(tvm_model, target="llvm")
+        vm = relax.VirtualMachine(ex, tvm_tirx.cpu())
+    vm.set_input("main", xs)
+    vm.invoke_stateful("main")
+    tvm_output = vm.get_outputs("main").numpy()
+
+    assert tvm_output.dtype == np_dtype
+    np.testing.assert_array_equal(tvm_output, expected)
+
+
 def test_leakyrelu():
     verify_unary("LeakyRelu", [32, 32])
     verify_unary("LeakyRelu", [32, 32], attrs={"alpha": 0.2})