[CINN] fix multi subgraph unittest (#61909)

BiynXu · web-flow · commit 0b70bbea095d · 2024-02-22T09:12:32.000+08:00
diff --git a/test/ir/pir/cinn/symbolic/test_llama_mlp_dy.py b/test/ir/pir/cinn/symbolic/test_llama_mlp_dy.py
@@ -34,13 +34,28 @@ def __init__(self):
         self.hidden_size = 768
         self.intermediate_size = 1008
         self.gate_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size, bias_attr=False
+            self.hidden_size,
+            self.intermediate_size,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Constant(value=0.5)
+            ),
+            bias_attr=False,
         )
         self.up_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size, bias_attr=False
+            self.hidden_size,
+            self.intermediate_size,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Constant(value=0.5)
+            ),
+            bias_attr=False,
         )
         self.down_proj = nn.Linear(
-            self.intermediate_size, self.hidden_size, bias_attr=False
+            self.intermediate_size,
+            self.hidden_size,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Constant(value=0.5)
+            ),
+            bias_attr=False,
         )
 
     def forward(self, x):
diff --git a/test/ir/pir/cinn/symbolic/test_llama_mlp_st.py b/test/ir/pir/cinn/symbolic/test_llama_mlp_st.py
@@ -34,13 +34,28 @@ def __init__(self):
         self.hidden_size = 768
         self.intermediate_size = 1008
         self.gate_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size, bias_attr=False
+            self.hidden_size,
+            self.intermediate_size,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Constant(value=0.5)
+            ),
+            bias_attr=False,
         )
         self.up_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size, bias_attr=False
+            self.hidden_size,
+            self.intermediate_size,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Constant(value=0.5)
+            ),
+            bias_attr=False,
         )
         self.down_proj = nn.Linear(
-            self.intermediate_size, self.hidden_size, bias_attr=False
+            self.intermediate_size,
+            self.hidden_size,
+            weight_attr=paddle.ParamAttr(
+                initializer=nn.initializer.Constant(value=0.5)
+            ),
+            bias_attr=False,
         )
 
     def forward(self, x):
diff --git a/test/ir/pir/cinn/symbolic/test_multiple_subgraph_dy.py b/test/ir/pir/cinn/symbolic/test_multiple_subgraph_dy.py
@@ -31,7 +31,12 @@ class MultipleSubgraph(nn.Layer):
     def __init__(self):
         super().__init__()
         self.hidden_size = 768
-        self.mlp = nn.Linear(self.hidden_size, self.hidden_size)
+        self.weight_attr = paddle.ParamAttr(
+            initializer=nn.initializer.Constant(value=0.5)
+        )
+        self.mlp = nn.Linear(
+            self.hidden_size, self.hidden_size, weight_attr=self.weight_attr
+        )
 
     def exp_sub(self, x):
         y = paddle.exp(x)
diff --git a/test/ir/pir/cinn/symbolic/test_multiple_subgraph_st.py b/test/ir/pir/cinn/symbolic/test_multiple_subgraph_st.py
@@ -31,7 +31,12 @@ class MultipleSubgraph(nn.Layer):
     def __init__(self):
         super().__init__()
         self.hidden_size = 768
-        self.mlp = nn.Linear(self.hidden_size, self.hidden_size)
+        self.weight_attr = paddle.ParamAttr(
+            initializer=nn.initializer.Constant(value=0.5)
+        )
+        self.mlp = nn.Linear(
+            self.hidden_size, self.hidden_size, weight_attr=self.weight_attr
+        )
 
     def exp_sub(self, x):
         y = paddle.exp(x)
