0d tensor for scatter_ and scatter_nd (#49072)

paddle/phi/infermeta/ternary.cc

@@ -1039,6 +1039,20 @@ void ScatterNdAddInferMeta(const MetaTensor& x,
   const auto& updates_dims = updates.dims();
   auto updates_dims_size = updates_dims.size();
 
+  if (updates_dims_size == 0) {
+    // check for 0d updates
+    PADDLE_ENFORCE_EQ(
+        index_dims_size,
+        1,
+        phi::errors::InvalidArgument("When the updates is a 0d tensor, the "
+                                     "index should be a 1d tensor."));
+    PADDLE_ENFORCE_EQ(
+        index_dims[index_dims_size - 1],
+        ref_dims_size,
+        phi::errors::InvalidArgument(
+            "When the update is a 0d tensor, The last dimension of "
+            "Input(Index)'s shape should be equal with the rank of Input(X)."));
+  } else {
     PADDLE_ENFORCE_LE(
         index_dims[index_dims_size - 1],
         ref_dims_size,
@@ -1063,7 +1077,7 @@ void ScatterNdAddInferMeta(const MetaTensor& x,
     for (int64_t i = index_dims[index_dims_size - 1]; i < ref_dims_size; ++i) {
       r_updates_dims.emplace_back(ref_dims[i]);
     }
-
+    // check for non-0d updates
     PADDLE_ENFORCE_EQ(
         r_updates_dims.size(),
         updates_dims_size,
@@ -1088,6 +1102,7 @@ void ScatterNdAddInferMeta(const MetaTensor& x,
               i,
               updates_dims[i]));
     }
+  }
   out->set_dims(ref_dims);
   out->share_lod(x);
   out->set_dtype(x.dtype());

python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py

@@ -682,6 +682,36 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(x2.grad.shape, [])
         self.assertEqual(x3.grad.shape, [])
 
+    def test_scatter__1D(self):
+        x = paddle.to_tensor([1.0, 3.0, 5.0, 7.0, 9.0])
+        index = paddle.full([], 2, 'int64')
+        updates = paddle.full([], 4.0)
+        out = paddle.scatter_(x, index, updates)
+
+        self.assertEqual(out.numpy()[2], 4)
+
+    def test_scatter__XD(self):
+        x = paddle.to_tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        index = paddle.full([], 1, 'int64')
+        updates = paddle.to_tensor([1.0, 2.0, 3.0])
+        out = paddle.scatter_(x, index, updates)
+
+        for i in range(3):
+            self.assertEqual(out.numpy()[1][i], updates.numpy()[i])
+
+    def test_scatter_nd(self):
+        index = paddle.to_tensor([3], dtype="int64", stop_gradient=False)
+        updates = paddle.full([], 2, dtype='float32')
+        updates.stop_gradient = False
+        shape = [5]
+
+        out = paddle.scatter_nd(index, updates, shape)
+        out.backward()
+
+        self.assertEqual(out.shape, [5])
+        self.assertEqual(out.numpy()[3], 2)
+        self.assertEqual(out.grad.shape, [5])
+
 
 class TestSundryAPIStatic(unittest.TestCase):
     def setUp(self):
@@ -845,6 +875,45 @@ class TestSundryAPIStatic(unittest.TestCase):
         self.assertEqual(res2.shape, (2, 2))
         self.assertEqual(res3.shape, (1, 1))
 
+    @prog_scope()
+    def test_scatter__1D(self):
+        x = paddle.full([10], 1.0, 'float32')
+        index = paddle.full([], 2, 'int64')
+        updates = paddle.full([], 4, 'float32')
+        out = paddle.scatter_(x, index, updates)
+        paddle.static.append_backward(out)
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, fetch_list=[out])
+        self.assertEqual(res[0][2], 4)
+
+    @prog_scope()
+    def test_scatter__XD(self):
+        x = paddle.full([2, 3], 1.0, 'float32')
+        index = paddle.full([], 1, 'int64')
+        updates = paddle.full([3], 4, 'float32')
+        out = paddle.scatter_(x, index, updates)
+        paddle.static.append_backward(out)
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, fetch_list=[out])
+        for i in range(3):
+            self.assertEqual(res[0][1][i], 4)
+
+    @prog_scope()
+    def test_scatter_nd(self):
+        index = paddle.static.data(name='index', shape=[1], dtype='int64')
+        updates = paddle.full([], 2, 'float32')
+        shape = [5]
+        index_data = np.array([3], dtype=np.longlong)
+        out = paddle.scatter_nd(index, updates, shape)
+        paddle.static.append_backward(out)
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, feed={'index': index_data}, fetch_list=[out])
+        self.assertEqual(res[0].shape, (5,))
+        self.assertEqual(res[0][3], 2)
+
 
 # Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.
 class TestNoBackwardAPI(unittest.TestCase):

python/paddle/fluid/tests/unittests/xpu/test_zero_dim_tensor_xpu.py

@@ -504,6 +504,23 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(x2.grad.shape, [])
         self.assertEqual(x3.grad.shape, [])
 
+    def test_scatter__1D(self):
+        x = paddle.to_tensor([1.0, 3.0, 5.0, 7.0, 9.0])
+        index = paddle.full([], 2, 'int64')
+        updates = paddle.full([], 4.0)
+        out = paddle.scatter_(x, index, updates)
+
+        self.assertEqual(out.numpy()[2], 4)
+
+    def test_scatter__XD(self):
+        x = paddle.to_tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        index = paddle.full([], 1, 'int64')
+        updates = paddle.to_tensor([1.0, 2.0, 3.0])
+        out = paddle.scatter_(x, index, updates)
+
+        for i in range(3):
+            self.assertEqual(out.numpy()[1][i], updates.numpy()[i])
+
 
 # Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.
 class TestNoBackwardAPI(unittest.TestCase):

python/paddle/tensor/manipulation.py

@@ -3078,7 +3078,7 @@ def scatter_nd(index, updates, shape, name=None):
     seen :code:`scatter_nd_add` . This op is the inverse of the :code:`gather_nd` op.
 
     Args:
-        index (Tensor): The index input with ndim > 1 and index.shape[-1] <= len(shape).
+        index (Tensor): The index input with ndim >= 1 and index.shape[-1] <= len(shape).
                           Its dtype should be int32 or int64 as it is used as indexes.
         updates (Tensor): The updated value of scatter_nd op. Its dtype should be float32, float64.
                             It must have the shape index.shape[:-1] + shape[index.shape[-1]:]