sort axes (apache#10985)

Co-authored-by: Margaret Qian <mqian@octoml.ai>
Lucien0 · Apr 19, 2022 · 7eb109e · 7eb109e
1 parent fd26465
commit 7eb109e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 5 deletions.
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
@@ -1505,7 +1505,8 @@ def _impl_v13(cls, inputs, attr, params):
 
         rank_input = len(infer_type(inputs[0]).checked_type.shape)
         num_new_axis = int(infer_type(inputs[1]).checked_type.shape[0])
-        axes = relay.split(inputs[1], num_new_axis).astuple()
+        axes = relay.sort(inputs[1])
+        axes = relay.split(axes, num_new_axis).astuple()
         result = inputs[0]
 
         # TODO (AndrewZhaoLuo): investigate performance issues with consecutive

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
@@ -5125,10 +5125,6 @@ def verify_eyelike(indata):
     "test_triu_square",
     "test_triu_square_neg",
     "test_triu_zero",
-    # These unsqueeze tests work, but take 2+ hrs to run
-    "test_unsqueeze_three_axes",
-    "test_unsqueeze_two_axes",
-    "test_unsqueeze_unsorted_axes",
     "test_unique_sorted_with_axis",
     "test_unique_sorted_with_axis_3d",
     "test_unique_sorted_with_negative_axis",