[Relax] Fix HardSigmoid returns 1.0 for NaN input

Author: cchung100m

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

@@ -3255,7 +3255,18 @@ def _impl_v1(cls, bb, inputs, attr, params):
         alpha = relax.const(alpha, dtype=dtype)
         beta = float(attr.get("beta", 0.5))
         beta = relax.const(beta, dtype=dtype)
-        return relax.op.clip(relax.op.add(relax.op.multiply(alpha, x), beta), 0, 1)
+        clipped = relax.op.clip(relax.op.add(relax.op.multiply(alpha, x), beta), 0, 1)
+
+        # Preserve NaN values: where x is NaN, return NaN instead of clipped value
+        if isinstance(x, relax.Constant):
+            x_data = x.data.numpy()
+            is_nan_data = _np.isnan(x_data)
+            is_nan = relax.const(is_nan_data, dtype="bool")
+        else:
+            is_nan = relax.op.not_equal(x, x)
+
+        nan_val = relax.const(_np.nan, dtype=dtype)
+        return relax.op.where(is_nan, nan_val, clipped)
 
 
 class HardSwish(OnnxOpConverter):

tests/python/relax/test_frontend_onnx.py

@@ -1086,6 +1086,31 @@ def test_hardsigmoid():
     verify_unary("HardSigmoid", [1, 3, 20, 20], attrs={"alpha": 0.5, "beta": 0.6})
 
 
+def test_hardsigmoid_nan():
+    """Test that HardSigmoid preserves NaN values in output."""
+    test_node = helper.make_node("HardSigmoid", ["x"], ["y"])
+    graph = helper.make_graph(
+        [test_node],
+        "hardsigmoid_nan_test",
+        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [3, 4])],
+        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 4])],
+    )
+
+    model = helper.make_model(graph, producer_name="hardsigmoid_nan_test")
+
+    # Create input with NaN values
+    input_data = np.array(
+        [
+            [np.nan, 0.5, -0.5, 1.0],
+            [0.0, np.nan, 2.0, -2.0],
+            [0.3, 0.7, np.nan, np.nan],
+        ],
+        dtype=np.float32,
+    )
+
+    check_correctness(model, inputs={"x": input_data})
+
+
 def test_shrink():
     verify_unary("Shrink", [32, 32])
     verify_unary("Shrink", [32, 32], attrs={"lambd": 0.2, "bias": 0.1})