[Paddle-TRT] fix expand_as_v2 conver bug (#67274)

ming1753 · web-flow · commit ccf97f0ab1e3 · 2024-08-14T16:30:56.000+08:00
* [Paddle-TRT] fix expand_as_v2 conver bug
diff --git a/paddle/fluid/inference/tensorrt/convert/expand_v2_op.cc b/paddle/fluid/inference/tensorrt/convert/expand_v2_op.cc
@@ -56,8 +56,9 @@ class ExpandOpConverter : public OpConverter {
       }
     } else if (op_type_ == "expand_as_v2") {
       if (inputs.find("Y") != inputs.end()) {
-        shape_tensor = engine_->GetITensor(op_desc.Input("Y")[0]);
-        shape_rank = shape_tensor->getDimensions().nbDims;
+        auto Y_t = engine_->GetITensor(op_desc.Input("Y")[0]);
+        shape_tensor = Shape(Y_t);
+        shape_rank = Y_t->getDimensions().nbDims;
       } else {
         std::vector<int32_t> shape = PADDLE_GET_CONST(
             std::vector<int32_t>, op_desc.GetAttr("target_shape"));
diff --git a/test/ir/inference/test_trt_convert_expand_v2.py b/test/ir/inference/test_trt_convert_expand_v2.py
@@ -415,5 +415,87 @@ def test(self):
         self.run_test()
 
 
+class TrtConvertExpandV2Test4(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+        return True
+
+    def sample_program_configs(self):
+        def generate_input1():
+            return np.random.random([1, 4, 1]).astype(np.float32)
+
+        def generate_input2():
+            return np.random.random([1, 4, 10]).astype(np.float32)
+
+        ops_config = [
+            {
+                "op_type": "expand_as_v2",
+                "op_inputs": {
+                    "X": ["expand_v2_input1"],
+                    "Y": ["expand_v2_input2"],
+                },
+                "op_outputs": {"Out": ["expand_v2_out"]},
+                "op_attrs": {"target_shape": [1, 4, 10]},
+            },
+        ]
+        ops = self.generate_op_config(ops_config)
+        program_config = ProgramConfig(
+            ops=ops,
+            weights={},
+            inputs={
+                "expand_v2_input1": TensorConfig(
+                    data_gen=partial(generate_input1)
+                ),
+                "expand_v2_input2": TensorConfig(
+                    data_gen=partial(generate_input2)
+                ),
+            },
+            outputs=["expand_v2_out"],
+        )
+
+        yield program_config
+
+    def sample_predictor_configs(
+        self, program_config
+    ) -> tuple[paddle_infer.Config, list[int], float]:
+        def generate_dynamic_shape():
+            self.dynamic_shape.min_input_shape = {
+                "expand_v2_input1": [1, 4, 1],
+                "expand_v2_input2": [1, 4, 10],
+            }
+            self.dynamic_shape.max_input_shape = {
+                "expand_v2_input1": [1, 4, 1],
+                "expand_v2_input2": [1, 4, 10],
+            }
+            self.dynamic_shape.opt_input_shape = {
+                "expand_v2_input1": [1, 4, 1],
+                "expand_v2_input2": [1, 4, 10],
+            }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.opt_input_shape = {}
+
+        clear_dynamic_shape()
+        # for dynamic_shape
+        generate_dynamic_shape()
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        program_config.set_input_type(np.float32)
+        yield self.create_inference_config(), (1, 3), 1e-5
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        program_config.set_input_type(np.float16)
+        yield self.create_inference_config(), (1, 3), 1e-3
+
+    def add_skip_trt_case(self):
+        pass
+
+    def test(self):
+        self.add_skip_trt_case()
+        self.run_test()
+
+
 if __name__ == "__main__":
     unittest.main()
