From ddc4246a7307b48e57ae1e7cbb2c94e9896a4b9c Mon Sep 17 00:00:00 2001
From: "Guan-Ming (Wesley) Chiu"
 <105915352+guan404ming@users.noreply.github.com>
Date: Fri, 19 Jun 2026 23:27:55 +0800
Subject: [PATCH] [Relax][ONNX] Accept 1-D scalar inputs in NonMaxSuppression

---
 .../tvm/relax/frontend/onnx/onnx_frontend.py  | 26 ++++++++++++++-----
 tests/python/relax/test_frontend_onnx.py      | 26 +++++++++++++++++++
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 3cfe7c892c46..739c5ce585f6 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -4913,7 +4913,7 @@ def _impl_v10(cls, bb, inputs, attr, params):
         if max_output_boxes_per_class is not None and isinstance(
             max_output_boxes_per_class, relax.Constant
         ):
-            max_output_boxes_per_class = int(max_output_boxes_per_class.data.numpy())
+            max_output_boxes_per_class = int(max_output_boxes_per_class.data.numpy().item())
         elif max_output_boxes_per_class is not None and isinstance(
             max_output_boxes_per_class, relax.Var
         ):
@@ -4927,12 +4927,19 @@ def _impl_v10(cls, bb, inputs, attr, params):
             max_output_boxes_per_class = 0  # Default value
 
         if iou_threshold is not None and isinstance(iou_threshold, relax.Constant):
-            iou_threshold = float(iou_threshold.data.numpy())
+            iou_threshold = float(iou_threshold.data.numpy().item())
+        elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
+            var_name = iou_threshold.name_hint
+            if var_name in params[1]:
+                _, param_value = params[1][var_name]
+                iou_threshold = float(param_value.numpy().item())
+            else:
+                iou_threshold = 0.5  # Default value
         else:
             iou_threshold = 0.5  # Default value
 
         if score_threshold is not None and isinstance(score_threshold, relax.Constant):
-            score_threshold = float(score_threshold.data.numpy())
+            score_threshold = float(score_threshold.data.numpy().item())
         elif score_threshold is not None and isinstance(score_threshold, relax.Var):
             var_name = score_threshold.name_hint
             if var_name in params[1]:
@@ -5002,7 +5009,7 @@ def _impl_v1(cls, bb, inputs, attr, params):
         if max_output_boxes_per_class is not None and isinstance(
             max_output_boxes_per_class, relax.Constant
         ):
-            max_output_boxes_per_class = int(max_output_boxes_per_class.data.numpy())
+            max_output_boxes_per_class = int(max_output_boxes_per_class.data.numpy().item())
         elif max_output_boxes_per_class is not None and isinstance(
             max_output_boxes_per_class, relax.Var
         ):
@@ -5016,12 +5023,19 @@ def _impl_v1(cls, bb, inputs, attr, params):
             max_output_boxes_per_class = 0  # Default value
 
         if iou_threshold is not None and isinstance(iou_threshold, relax.Constant):
-            iou_threshold = float(iou_threshold.data.numpy())
+            iou_threshold = float(iou_threshold.data.numpy().item())
+        elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
+            var_name = iou_threshold.name_hint
+            if var_name in params[1]:
+                _, param_value = params[1][var_name]
+                iou_threshold = float(param_value.numpy().item())
+            else:
+                iou_threshold = 0.5  # Default value
         else:
             iou_threshold = 0.5  # Default value
 
         if score_threshold is not None and isinstance(score_threshold, relax.Constant):
-            score_threshold = float(score_threshold.data.numpy())
+            score_threshold = float(score_threshold.data.numpy().item())
         elif score_threshold is not None and isinstance(score_threshold, relax.Var):
             var_name = score_threshold.name_hint
             if var_name in params[1]:
diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py
index 57f780868ccd..3de10cb25664 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -5190,6 +5190,32 @@ def test_nms():
         )
 
 
+def test_nms_scalar_shape1_constants():
+    """Scalar params given as 1-D single-element constants must import (NumPy 2.x cast)."""
+    nms_node = helper.make_node(
+        "NonMaxSuppression",
+        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
+        ["selected_indices"],
+    )
+    graph = helper.make_graph(
+        [nms_node],
+        "nms_scalar_shape1",
+        inputs=[
+            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, [1, 5, 4]),
+            helper.make_tensor_value_info("scores", TensorProto.FLOAT, [1, 1, 5]),
+        ],
+        initializer=[
+            helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [3]),
+            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.5]),
+            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.0]),
+        ],
+        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [0, 3])],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
+    # Default import folds initializers to relax.Constant, exercising the scalar-cast path.
+    from_onnx(model)
+
+
 @pytest.mark.parametrize("with_explicit_max", [False, True])
 def test_nms_max_output_boxes_per_class_zero(with_explicit_max: bool):
     """ONNX default for max_output_boxes_per_class is 0, yielding empty output."""