From 414b60f461c5b70121669684c3933f29ed20c134 Mon Sep 17 00:00:00 2001 From: leeleolay <805351297@qq.com> Date: Thu, 16 Jul 2026 23:52:00 +0800 Subject: [PATCH] Add scatter_nd_add double gradient support --- .../op_generator/vjp_interface_black_list.py | 1 + .../composite_double_backward_api.h | 25 +++++++ paddle/phi/ops/yaml/backward.yaml | 11 +++ test/legacy_test/test_scatter_nd_op.py | 74 +++++++++++++++++++ 4 files changed, 111 insertions(+) diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py index fd9d770f4e2b8..0035fe582219b 100644 --- a/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py +++ b/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py @@ -39,6 +39,7 @@ 'index_add_grad', 'acos_grad', 'put_along_axis_grad', + 'scatter_nd_add_grad', 'masked_fill_grad', 'masked_select_grad', 'index_elementwise_get_grad', diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h index 33a86eceb7000..4e3fc87cce709 100644 --- a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h +++ b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h @@ -1259,6 +1259,31 @@ void put_along_axis_double_grad(const Tensor& arr, } } +template +void scatter_nd_add_double_grad(const Tensor& index, + const Tensor& grad_out, + const optional& grad_x_grad, + const optional& grad_updates_grad, + Tensor* grad_out_grad) { + if (grad_out_grad) { + Tensor grad_out_grad_tmp; + if (grad_x_grad) { + grad_out_grad_tmp = grad_x_grad.get(); + } else { + grad_out_grad_tmp = full(common::vectorize(grad_out.dims()), + 0, + grad_out.dtype(), + grad_out.place()); + } + + if (grad_updates_grad) { + grad_out_grad_tmp = + scatter_nd_add(grad_out_grad_tmp, index, grad_updates_grad.get()); + } + set_output(grad_out_grad_tmp, grad_out_grad); + } +} + template void index_add_double_grad(const Tensor& index, const Tensor& out_grad, diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml index 4d3b2fb19c78b..6042282ac1b54 100644 --- a/paddle/phi/ops/yaml/backward.yaml +++ b/paddle/phi/ops/yaml/backward.yaml @@ -3278,6 +3278,16 @@ no_need_buffer : updates composite: scatter_grad(index, updates, out_grad, overwrite, x_grad, updates_grad) +- backward_op : scatter_nd_add_double_grad + forward : scatter_nd_add_grad (Tensor index, Tensor updates, Tensor grad_out) -> Tensor(grad_x), Tensor(grad_updates) + args : (Tensor index, Tensor grad_out, Tensor grad_x_grad, Tensor grad_updates_grad) + output : Tensor(grad_out_grad) + infer_meta : + func : UnchangedInferMeta + param : [grad_out] + composite : scatter_nd_add_double_grad(index, grad_out, grad_x_grad, grad_updates_grad, grad_out_grad) + optional : grad_x_grad, grad_updates_grad + - backward_op : scatter_nd_add_grad forward : scatter_nd_add (Tensor x, Tensor index, Tensor updates) -> Tensor(out) args : (Tensor index, Tensor updates, Tensor out_grad) @@ -3289,6 +3299,7 @@ func : scatter_nd_add_grad no_need_buffer : updates composite: scatter_nd_add_grad(index, updates, out_grad, x_grad, updates_grad) + backward : scatter_nd_add_double_grad - backward_op : segment_pool_grad forward : segment_pool (Tensor x, Tensor segment_ids, str pooltype="SUM") -> Tensor(out), Tensor(summed_ids) diff --git a/test/legacy_test/test_scatter_nd_op.py b/test/legacy_test/test_scatter_nd_op.py index a470daba7a1d2..64062d4516763 100644 --- a/test/legacy_test/test_scatter_nd_op.py +++ b/test/legacy_test/test_scatter_nd_op.py @@ -645,6 +645,80 @@ def test_dygraph(self): ) +class TestScatterNdAddDoubleGrad(unittest.TestCase): + def _compute_double_grad(self, first_grad): + x = paddle.to_tensor([[0.2, 0.4], [-0.5, 0.1]], dtype="float64") + updates = paddle.to_tensor( + [[0.7, -0.3], [0.5, 0.2], [-0.4, 0.8]], + dtype="float64", + ) + index = paddle.to_tensor([[0], [0], [1]], dtype="int64") + x.stop_gradient = False + updates.stop_gradient = False + + output = paddle.scatter_nd_add(x, index, updates) + loss = output.square().sum() + x_grad, updates_grad = paddle.grad( + loss, [x, updates], create_graph=True + ) + if first_grad == "x": + grad_sum = x_grad.sum() + elif first_grad == "updates": + grad_sum = updates_grad.sum() + else: + grad_sum = x_grad.sum() + updates_grad.sum() + return paddle.grad(grad_sum, [x, updates]) + + def test_create_graph(self): + with base.dygraph.guard(base.CPUPlace()): + updates = paddle.randn([3, 2], dtype="float64") + updates.stop_gradient = False + index = paddle.to_tensor([[0], [0], [1]], dtype="int64") + output = paddle.scatter_nd_add( + paddle.zeros([2, 2], dtype="float64"), index, updates + ) + (updates_grad,) = paddle.grad( + output.square().sum(), updates, create_graph=True + ) + self.assertEqual(updates_grad.shape, updates.shape) + + def test_double_grad(self): + cases = ( + ( + "x", + [[2.0, 2.0], [2.0, 2.0]], + [[2.0, 2.0], [2.0, 2.0], [2.0, 2.0]], + ), + ( + "updates", + [[4.0, 4.0], [2.0, 2.0]], + [[4.0, 4.0], [4.0, 4.0], [2.0, 2.0]], + ), + ( + "both", + [[6.0, 6.0], [4.0, 4.0]], + [[6.0, 6.0], [6.0, 6.0], [4.0, 4.0]], + ), + ) + original_prim = core._is_eager_prim_enabled() + core.set_prim_eager_enabled(True) + try: + with base.dygraph.guard(base.CPUPlace()): + for first_grad, expected_x, expected_updates in cases: + with self.subTest(first_grad=first_grad): + x_double_grad, updates_double_grad = ( + self._compute_double_grad(first_grad) + ) + np.testing.assert_allclose( + x_double_grad.numpy(), expected_x + ) + np.testing.assert_allclose( + updates_double_grad.numpy(), expected_updates + ) + finally: + core.set_prim_eager_enabled(original_prim) + + if __name__ == "__main__": paddle.enable_static() unittest.main()