[PIR]Fix array_write stop_gradient (#60970)

changeyoung98 · web-flow · commit 4db394f5a530 · 2024-01-22T23:23:29.000+08:00
* fix array_write stop_gradient

* fix stop_gradient bug

* set array_write sg to true
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -2179,7 +2179,14 @@ void ArrayWrite_Op::Build(pir::Builder &builder,
       ArrayWrite_Op::InferMeta(argument_inputs, argument_attributes);
 
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
-  ::pir::PassStopGradientsDefaultly(argument);
+  constexpr char kStopGradientAttrName[] = "stop_gradient";
+  auto stop_gradient0 =
+      argument.inputs[0].attribute<pir::BoolAttribute>(kStopGradientAttrName);
+  auto stop_gradient1 =
+      argument.inputs[1].attribute<pir::BoolAttribute>(kStopGradientAttrName);
+  auto stop_gradient = stop_gradient0.data() && stop_gradient1.data();
+  argument.inputs[0].set_attribute(kStopGradientAttrName,
+                                   builder.bool_attr(stop_gradient));
 }
 
 void ArrayWrite_Op::VerifySig() {
diff --git a/test/legacy_test/test_assign_op.py b/test/legacy_test/test_assign_op.py
@@ -129,8 +129,6 @@ def test_assign_tensor_array(self):
             z = paddle.add(x=x, y=y)
             i = paddle.tensor.fill_constant(shape=[1], dtype='int64', value=0)
             init_array = paddle.tensor.array_write(x=z, i=i)
-            # TODO(xiaoguoguo626807): Remove this stop_gradient=False.
-            init_array.stop_gradient = False
             array = paddle.assign(init_array)
             sums = paddle.tensor.array_read(array=init_array, i=i)
             mean = paddle.mean(sums)
diff --git a/test/legacy_test/test_dynamic_rnn_stop_gradient.py b/test/legacy_test/test_dynamic_rnn_stop_gradient.py
@@ -59,6 +59,7 @@ def build_and_run_program(place, batch_size, beam_size, stop_gradient=False):
         paddle.tensor.array_write(score, i=step_idx, array=scores)
         length_cond = paddle.less_than(x=step_idx, y=max_len)
         paddle.assign(length_cond, cond)
+    scores.stop_gradient = True
     out = tensor_array_to_tensor(scores, axis=0, use_stack=True)[0]
     loss = paddle.mean(out)
     opt = paddle.optimizer.Adam(0.01)
