From 3dcb490b5798f93f44a45363f361a9a62b8677ce Mon Sep 17 00:00:00 2001 From: Margaret Qian Date: Tue, 12 Apr 2022 14:46:17 -0700 Subject: [PATCH] sort axes --- python/tvm/relay/frontend/onnx.py | 3 ++- tests/python/frontend/onnx/test_forward.py | 4 ---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 7dcb9952c7fb..b2f8467704bc 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -1499,7 +1499,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 index 638b4b8f57eb..24600f49caa6 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -5132,10 +5132,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",