From cb346e344cfcb391e808f2721731c2558be3896d Mon Sep 17 00:00:00 2001
From: Gaoyuyang <zdlnlhmj@126.com>
Date: Thu, 2 Mar 2023 08:52:10 +0000
Subject: [PATCH] add take_alone_axis and topk converter for paddle frontend

---
 python/tvm/relay/frontend/paddlepaddle.py     | 18 +++++++++++++--
 .../frontend/paddlepaddle/test_forward.py     | 23 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relay/frontend/paddlepaddle.py b/python/tvm/relay/frontend/paddlepaddle.py
index 78895e4b49e0..162a850489cc 100755
--- a/python/tvm/relay/frontend/paddlepaddle.py
+++ b/python/tvm/relay/frontend/paddlepaddle.py
@@ -2025,6 +2025,14 @@ def convert_swish(g, op, block):
     g.add_node(op.output("Out")[0], out)
 
 
+def convert_take_along_axis(g, op, block):
+    x = g.get_node(op.input("Input")[0])
+    idx = g.get_node(op.input("Index")[0])
+    axis = op.attr("Axis")
+    out = _op.gather(x, axis, idx)
+    g.add_node(op.output("Result")[0], out)
+
+
 def convert_topk(g, op, block):
     """Operator converter for topk."""
 
@@ -2034,9 +2042,13 @@ def convert_topk(g, op, block):
     else:
         k = op.attr("k")
 
-    largest = op.attr("largest")
+    largest = True
+    axis = -1
+    if op.has_attr("axis"):
+        axis = op.attr("axis")
+    if op.has_attr("largest"):
+        largest = op.attr("largest")
     is_ascend = not largest
-    axis = op.attr("axis")
 
     value_names = op.output("Out")
     indice_names = op.output("Indices")
@@ -2214,8 +2226,10 @@ def convert_where_index(g, op, block):
     "square": convert_square,
     "squeeze2": convert_squeeze,
     "swish": convert_swish,
+    "take_along_axis": convert_take_along_axis,
     "tan": convert_unary_op,
     "tanh": convert_unary_op,
+    "top_k": convert_topk,
     "top_k_v2": convert_topk,
     "transpose2": convert_transpose,
     "unsqueeze2": convert_unsqueeze,
diff --git a/tests/python/frontend/paddlepaddle/test_forward.py b/tests/python/frontend/paddlepaddle/test_forward.py
index 70fbf6aee554..8d501c0ca52f 100755
--- a/tests/python/frontend/paddlepaddle/test_forward.py
+++ b/tests/python/frontend/paddlepaddle/test_forward.py
@@ -1739,13 +1739,25 @@ def topk5(inputs):
     def topk6(inputs):
         return paddle.topk(inputs, k=1, axis=0)
 
+    # paddle.fluid.layers.topk
+    @paddle.jit.to_static
+    def topk7(inputs):
+        return paddle.fluid.layers.topk(inputs, k=1)
+
+    @paddle.jit.to_static
+    def topk8(inputs):
+        return paddle.fluid.layers.topk(inputs, k=2)
+
     input_data = paddle.to_tensor([[1, 4, 5, 7], [3, 6, 2, 5]], dtype=paddle.int32)
+    input_data_fp32 = paddle.to_tensor([[1, 4, 5, 7], [3, 6, 2, 5]], dtype=paddle.float32)
     verify_model(topk1, input_data=input_data)
     # verify_model(topk2, input_data=input_data)
     verify_model(topk3, input_data=input_data)
     verify_model(topk4, input_data=input_data)
     verify_model(topk5, input_data=input_data)
     verify_model(topk6, input_data=input_data)
+    verify_model(topk7, input_data=input_data_fp32)
+    verify_model(topk8, input_data=input_data_fp32)
 
 
 @tvm.testing.uses_gpu
@@ -1784,5 +1796,16 @@ def where_index_1(inputs):
     verify_model(where_index_1, input_data=input_data, use_vm=True)
 
 
+@tvm.testing.uses_gpu
+def test_forward_take_along_axis():
+    @paddle.jit.to_static
+    def take_along_axis_1(inputs, index):
+        return paddle.take_along_axis(inputs, index, 0)
+
+    input_data = paddle.to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    index = paddle.to_tensor([[0]])
+    verify_model(take_along_axis_1, input_data=[input_data, index])
+
+
 if __name__ == "__main__":
     tvm.testing.main()