Specifying input data types

PaddlePaddle · May 19, 2023 · f15b894 · f15b894
1 parent ab770f8
commit f15b894
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 7 deletions.
diff --git a/test/ir/inference/auto_scan_test.py b/test/ir/inference/auto_scan_test.py
@@ -763,6 +763,7 @@ def random_to_skip():
                 threshold,
             ) in self.sample_predictor_configs(prog_config):
                 for input_type in self.get_avalible_input_type():
+                    prog_config = prog_config.set_input_type(input_type)
                     if os.path.exists(self.cache_dir):
                         shutil.rmtree(self.cache_dir)
 

diff --git a/test/ir/inference/test_multihead_matmul_fuse_pass_v3.py b/test/ir/inference/test_multihead_matmul_fuse_pass_v3.py
@@ -33,6 +33,9 @@ def generate_mul_input():
         def generate_elewise_input():
             return np.random.random([1, 12, 128, 128]).astype(np.float32)
 
+        def generate_weight(shape):
+            return np.random.random(shape).astype(np.float32)
+
         mul_0 = OpConfig(
             "mul",
             inputs={"X": ["mul_x"], "Y": ["mul_0_w"]},
@@ -195,13 +198,27 @@ def generate_elewise_input():
                 ),
             },
             weights={
-                "mul_0_w": TensorConfig(shape=[768, 768]),
-                "mul_1_w": TensorConfig(shape=[768, 768]),
-                "mul_2_w": TensorConfig(shape=[768, 768]),
-                "mul_3_w": TensorConfig(shape=[768, 768]),
-                "ele_0_w": TensorConfig(shape=[768]),
-                "ele_1_w": TensorConfig(shape=[768]),
-                "ele_2_w": TensorConfig(shape=[768]),
+                "mul_0_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768, 768])
+                ),
+                "mul_1_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768, 768])
+                ),
+                "mul_2_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768, 768])
+                ),
+                "mul_3_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768, 768])
+                ),
+                "ele_0_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768])
+                ),
+                "ele_1_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768])
+                ),
+                "ele_2_w": TensorConfig(
+                    data_gen=partial(generate_weight, [768])
+                ),
             },
             outputs=[ops[-1].outputs["Out"][0]],
         )