[Operator Mechanism] Add scatter_nd_add double gradient support

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',

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 <typename T>
+void scatter_nd_add_double_grad(const Tensor& index,
+                                const Tensor& grad_out,
+                                const optional<Tensor>& grad_x_grad,
+                                const optional<Tensor>& 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<T>(common::vectorize(grad_out.dims()),
+                                  0,
+                                  grad_out.dtype(),
+                                  grad_out.place());
+    }
+
+    if (grad_updates_grad) {
+      grad_out_grad_tmp =
+          scatter_nd_add<T>(grad_out_grad_tmp, index, grad_updates_grad.get());
+    }
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 template <typename T>
 void index_add_double_grad(const Tensor& index,
                            const Tensor& out_grad,

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)

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()