[tvm] branch main updated: [ONNX] Use take instead of min in NMS conditions (#7633)

[mb...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mbrookhart pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new f8596b5  [ONNX] Use take instead of min in NMS conditions (#7633)
f8596b5 is described below

commit f8596b5228346eae18f25f22e3096fef758b1166
Author: masahi <ma...@gmail.com>
AuthorDate: Fri Mar 12 00:56:28 2021 +0900

    [ONNX] Use take instead of min in NMS conditions (#7633)
---
 python/tvm/relay/frontend/onnx.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 860753d..f31b8c9 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -2453,7 +2453,7 @@ class NonMaxSuppression(OnnxOpConverter):
             nms_size_out,
         ):
             # Loop over classes, end when i == C
-            return _op.min(_op.less(i, C))
+            return _op.take(_op.less(i, C), _expr.const(0))
 
         def _first_body(
             i,
@@ -2561,7 +2561,7 @@ class NonMaxSuppression(OnnxOpConverter):
 
         def _inner_cond(i, j, C, onnx_out, nms_size, out):
             # inner loop over number of classes
-            return _op.min(_op.less(j, C))
+            return _op.take(_op.less(j, C), _expr.const(0))
 
         def _inner_body(i, j, C, onnx_out, nms_size, out):
             # slice to get current batch and class for valid box indicator
@@ -2591,7 +2591,7 @@ class NonMaxSuppression(OnnxOpConverter):
 
         def _outer_cond(i, B, C, onnx_out, nms_size_out, out):
             # Outer loop is over batch size
-            return _op.min(_op.less(i, B))
+            return _op.take(_op.less(i, B), _expr.const(0))
 
         def _outer_body(i, B, C, onnx_out, nms_size_out, out):
             # Outer loop just calls inner loop