【BUPT】【Paddle TensorRT No.38】Add pd_op.gather converter

python/paddle/tensorrt/impls/manipulation.py

@@ -70,6 +70,17 @@ def reshape_converter(network, paddle_op, inputs):
     return shuffle_layer.get_output(0)
 
 
+@converter_registry.register("pd_op.gather", trt_version="8.x")
+def gather_converter(network, paddle_op, inputs):
+    input_tensor, index_tensor, axis = inputs
+    reshape_layer = network.add_shuffle(index_tensor)
+    reshape_layer.reshape_dims = (-1,)
+    gather_layer = network.add_gather(
+        input_tensor, reshape_layer.get_output(0), axis
+    )
+    return gather_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.gather_nd", trt_version="8.x")
 def gather_nd_converter(network, paddle_op, inputs):
     input_tensor, indices_tensor = inputs

test/tensorrt/test_converter_manipulation.py

@@ -567,5 +567,37 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestGatherTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.gather
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("float32"),
+            "index": np.array([0, 2]).astype("int64"),
+            "axis": np.array([1]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "index", "axis"]}
+        self.min_shape = {"x": [1, 4, 10], "index": [1], "axis": [1]}
+        self.max_shape = {"x": [5, 4, 10], "index": [5], "axis": [1]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGatherCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.gather
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("int64"),
+            "index": np.array([0, 2]).astype("int64"),
+            "axis": np.array([1]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "index", "axis"]}
+        self.min_shape = {"x": [1, 4, 10], "index": [1], "axis": [1]}
+        self.max_shape = {"x": [5, 4, 10], "index": [5], "axis": [1]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()