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

python/paddle/tensorrt/impls/manipulation.py

@@ -70,6 +70,15 @@ 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, *_ = inputs
+    gather_layer = network.add_gather_v2(
+        input_tensor, index_tensor, trt.GatherMode.DEFAULT
+    )
+    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,25 @@ 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": 1,
+        }
+        self.program_config = {"feed_list": ["x", "index"]}
+        self.min_shape = {"x": [1, 4, 10], "index": [1]}
+        self.max_shape = {"x": [5, 4, 10], "index": [5]}
+        self.dynamic_shape = {
+            "x": {"min": [1, 4, 10], "max": [5, 4, 10], "opt": [3, 4, 10]},
+            "index": {"min": [1], "max": [5], "opt": [2]},
+        }
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()