From 4db394f5a530e9f1a324ca272fcbe4c442e5a747 Mon Sep 17 00:00:00 2001 From: Chen Zhiyang <1792266893@qq.com> Date: Mon, 22 Jan 2024 23:23:29 +0800 Subject: [PATCH] [PIR]Fix array_write stop_gradient (#60970) * fix array_write stop_gradient * fix stop_gradient bug * set array_write sg to true --- paddle/fluid/pir/dialect/operator/ir/manual_op.cc | 9 ++++++++- test/legacy_test/test_assign_op.py | 2 -- test/legacy_test/test_dynamic_rnn_stop_gradient.py | 1 + 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc index 497282e48909d..9e387d4ba85be 100644 --- 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(kStopGradientAttrName); + auto stop_gradient1 = + argument.inputs[1].attribute(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 index b4ccbf56a72d4..5a2ff191a712f 100644 --- 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 index c6d85b864c8d7..8ffaca6b456ec 100644 --- 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)