From 74582aaa190ea9fa0994687f3b9c4cf67c73e99d Mon Sep 17 00:00:00 2001 From: Yuang Liu Date: Fri, 16 Dec 2022 11:06:42 +0800 Subject: [PATCH] 0d tensor for scatter_ and scatter_nd (#49072) --- paddle/phi/infermeta/ternary.cc | 107 ++++++++++-------- .../tests/unittests/test_zero_dim_tensor.py | 69 +++++++++++ .../unittests/xpu/test_zero_dim_tensor_xpu.py | 17 +++ python/paddle/tensor/manipulation.py | 2 +- 4 files changed, 148 insertions(+), 47 deletions(-) diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc index f7bae369099..74391958189 100644 --- a/paddle/phi/infermeta/ternary.cc +++ b/paddle/phi/infermeta/ternary.cc @@ -1039,54 +1039,69 @@ void ScatterNdAddInferMeta(const MetaTensor& x, const auto& updates_dims = updates.dims(); auto updates_dims_size = updates_dims.size(); - PADDLE_ENFORCE_LE( - index_dims[index_dims_size - 1], - ref_dims_size, - phi::errors::InvalidArgument( - "The last dimension of Input(Index)'s shape should be no greater " - "than the rank of Input(X), but received the last dimension of " - "Input(Index)'s shape is %d, the rank of Input(X) is %d.", - index_dims[index_dims_size - 1], - ref_dims_size)); - PADDLE_ENFORCE_GE(index_dims_size, - 2UL, - phi::errors::InvalidArgument( - "The rank of Input(Index) should be greater than 1, " - "but received the rank of Input(Index) is %d.", - index_dims_size)); - - // update.shape = index.shape[:-1] + output.shape[index.shape[-1]:] - std::vector r_updates_dims; - for (int64_t i = 0; i < index_dims_size - 1; ++i) { - r_updates_dims.emplace_back(index_dims[i]); - } - for (int64_t i = index_dims[index_dims_size - 1]; i < ref_dims_size; ++i) { - r_updates_dims.emplace_back(ref_dims[i]); - } - - PADDLE_ENFORCE_EQ( - r_updates_dims.size(), - updates_dims_size, - phi::errors::InvalidArgument( - "Updates has wrong shape. The shape of Updates and Input(Updates) " - "should be same, but received the shape of Updates is %d, " - "the shape of Input(Updates) is %d.", - r_updates_dims.size(), - updates_dims_size)); - - for (int64_t i = 0; i < updates_dims_size; ++i) { + 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( - r_updates_dims[i], - updates_dims[i], + index_dims[index_dims_size - 1], + ref_dims_size, phi::errors::InvalidArgument( - "Updates has wrong shape. The dimensions of Updates and " - "Input(Updates) should match, but received Updates's" - "%d-th dimension is %d, Input(Updates)'s %d-th " - "dimension is %d.", - i, - r_updates_dims[i], - i, - updates_dims[i])); + "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, + phi::errors::InvalidArgument( + "The last dimension of Input(Index)'s shape should be no greater " + "than the rank of Input(X), but received the last dimension of " + "Input(Index)'s shape is %d, the rank of Input(X) is %d.", + index_dims[index_dims_size - 1], + ref_dims_size)); + PADDLE_ENFORCE_GE(index_dims_size, + 2UL, + phi::errors::InvalidArgument( + "The rank of Input(Index) should be greater than 1, " + "but received the rank of Input(Index) is %d.", + index_dims_size)); + + // update.shape = index.shape[:-1] + output.shape[index.shape[-1]:] + std::vector r_updates_dims; + for (int64_t i = 0; i < index_dims_size - 1; ++i) { + r_updates_dims.emplace_back(index_dims[i]); + } + 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, + phi::errors::InvalidArgument( + "Updates has wrong shape. The shape of Updates and Input(Updates) " + "should be same, but received the shape of Updates is %d, " + "the shape of Input(Updates) is %d.", + r_updates_dims.size(), + updates_dims_size)); + + for (int64_t i = 0; i < updates_dims_size; ++i) { + PADDLE_ENFORCE_EQ( + r_updates_dims[i], + updates_dims[i], + phi::errors::InvalidArgument( + "Updates has wrong shape. The dimensions of Updates and " + "Input(Updates) should match, but received Updates's" + "%d-th dimension is %d, Input(Updates)'s %d-th " + "dimension is %d.", + i, + r_updates_dims[i], + i, + updates_dims[i])); + } } out->set_dims(ref_dims); out->share_lod(x); diff --git a/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py b/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py index eae5528fba2..9fb6017d446 100644 --- a/python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py +++ b/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): diff --git a/python/paddle/fluid/tests/unittests/xpu/test_zero_dim_tensor_xpu.py b/python/paddle/fluid/tests/unittests/xpu/test_zero_dim_tensor_xpu.py index a6f91e5df4c..018ecc20e7d 100644 --- a/python/paddle/fluid/tests/unittests/xpu/test_zero_dim_tensor_xpu.py +++ b/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): diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py index df478ddde46..c255a223683 100644 --- a/python/paddle/tensor/manipulation.py +++ b/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]:] -- GitLab