add a shape function and dynamic test for round (apache#7324)

electriclilies · Feb 18, 2021 · c11aa3a · c11aa3a
1 parent f2fe840
commit c11aa3a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py
@@ -235,6 +235,7 @@ def elemwise_shape_func(attrs, inputs, _):
 
 register_shape_func("cast", False, elemwise_shape_func)
 register_shape_func("cast_like", False, elemwise_shape_func)
+register_shape_func("round", False, elemwise_shape_func)
 register_shape_func("zeros", False, no_data_full_shape_func)
 register_shape_func("zeros_like", False, elemwise_shape_func)
 register_shape_func("ones", False, no_data_full_shape_func)

diff --git a/tests/python/relay/test_any.py b/tests/python/relay/test_any.py
@@ -120,6 +120,7 @@ def test_any_elemwise():
     verify_any_elemwise((relay.Any(),), (3,), relay.sqrt, np.sqrt)
     verify_any_elemwise((relay.Any(), 2), (5, 2), relay.negative, np.negative)
     verify_any_elemwise((relay.Any(), relay.Any()), (5, 4), relay.exp, np.exp)
+    verify_any_elemwise((relay.Any(),), (3,), relay.round, np.round)
 
 
 @tvm.testing.uses_gpu