【pir】modify test_array_read_write.py (#65856)

xiaoguoguo626807 · web-flow · commit e500b47560d9 · 2024-07-10T11:06:13.000+08:00
* modify test_array_read_write.py

* modify test_array_read_write.py
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -327,7 +327,7 @@
   kernel :
     func : fetch
     param : [x]
-  traits : pir::SideEffectTrait, pir::ImmutableLayoutTrait
+  traits : pir::SideEffectTrait, pir::ImmutableLayoutTrait, paddle::dialect::ForwardOnlyTrait
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : floor_divide
diff --git a/test/legacy_test/test_array_read_write_op.py b/test/legacy_test/test_array_read_write_op.py
@@ -83,14 +83,21 @@ def test_read_write(self):
         total_sum = paddle.add_n([a_sum, x_sum])
         total_sum_scaled = paddle.scale(x=total_sum, scale=1 / 6.0)
 
-        append_backward(total_sum_scaled)
-
-        g_vars = list(
-            map(
-                default_main_program().global_block().var,
-                [each_x.name + "@GRAD" for each_x in x],
+        grad_list = append_backward(total_sum_scaled, [x[0], x[1], x[2]])
+        if not paddle.framework.in_pir_mode():
+            g_vars = list(
+                map(
+                    default_main_program().global_block().var,
+                    [each_x.name + "@GRAD" for each_x in x],
+                )
             )
-        )
+        else:
+            g_vars = []
+            for each_x in x:
+                for p, g in grad_list:
+                    if p.is_same(each_x):
+                        g_vars.append(g)
+                        continue
         g_out = [
             item.sum()
             for item in exe.run(
