Add register_shape_func for argsort (apache#6978)

neo-ai · Jun 12, 2022 · 8227062 · 8227062
1 parent d2755c6
commit 8227062
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/tvm/relay/op/_algorithm.py b/python/tvm/relay/op/_algorithm.py
@@ -24,11 +24,13 @@
 from . import strategy
 from . import op as _reg
 from .op import OpPattern, register_pattern
-from .op import register_strategy
+from .op import register_strategy, register_shape_func
+from ._tensor import elemwise_shape_func
 
 # argsort
 register_strategy("argsort", strategy.argsort_strategy)
 register_pattern("argsort", OpPattern.OPAQUE)
+register_shape_func("argsort", False, elemwise_shape_func)
 
 # topk
 register_strategy("topk", strategy.topk_strategy)