From 3a0aeaa41c501d764024b7ceb5a50d66c1701abc Mon Sep 17 00:00:00 2001 From: YibLiu <68105073+YibinLiu666@users.noreply.github.com> Date: Thu, 7 Dec 2023 10:22:53 +0800 Subject: [PATCH] =?UTF-8?q?=E3=80=90Hackathon=205th=20No.27=E3=80=91?= =?UTF-8?q?=E4=B8=BA=20Paddle=20=E6=96=B0=E5=A2=9E=20select=5Fscatter=20AP?= =?UTF-8?q?I=20-part=20(#59343)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * support select_scatter op * fix example code * fix sc * update example * remove unused files * add name * fix conflict * update * remove * update * add type * update type --- python/paddle/__init__.py | 2 + python/paddle/tensor/__init__.py | 2 + python/paddle/tensor/manipulation.py | 106 +++++++++++++ test/legacy_test/test_select_scatter_op.py | 167 +++++++++++++++++++++ 4 files changed, 277 insertions(+) create mode 100644 test/legacy_test/test_select_scatter_op.py diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py index df41a00202f2b..f6b3aca02c2a6 100644 --- a/python/paddle/__init__.py +++ b/python/paddle/__init__.py @@ -242,6 +242,7 @@ tolist, take_along_axis, put_along_axis, + select_scatter, tensordot, as_complex, as_real, @@ -894,6 +895,7 @@ 'renorm_', 'take_along_axis', 'put_along_axis', + 'select_scatter', 'multigammaln', 'multigammaln_', 'nan_to_num', diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py index dc69a4e58bb9a..1682378f59f90 100644 --- a/python/paddle/tensor/__init__.py +++ b/python/paddle/tensor/__init__.py @@ -175,6 +175,7 @@ scatter_, scatter_nd, scatter_nd_add, + select_scatter, shard_index, slice, split, @@ -697,6 +698,7 @@ 'repeat_interleave', 'take_along_axis', 'put_along_axis', + 'select_scatter', 'put_along_axis_', 'exponential_', 'heaviside', diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py index b60a12dd78370..91a3887d66387 100644 --- a/python/paddle/tensor/manipulation.py +++ b/python/paddle/tensor/manipulation.py @@ -6006,3 +6006,109 @@ def diagonal_scatter(x, y, offset=0, axis1=0, axis2=1, name=None): """ return fill_diagonal_tensor(x, y, offset, axis1, axis2, name) + + +def select_scatter(x, values, axis, index, name=None): + """ + Embeds the values of the values tensor into x at the given index of axis. + + Args: + x (Tensor) : The Destination Tensor. Supported data types are `bool`, `float16`, `float32`, `float64`, `uint8`, `int8`, `int16`, `int32`, `int64`, `bfloat16`, `complex64`, `complex128`. + values (Tensor) : The tensor to embed into x. Supported data types are `bool`, `float16`, `float32`, `float64`, `uint8`, `int8`, `int16`, `int32`, `int64`, `bfloat16`, `complex64`, `complex128`. + axis (int) : the dimension to insert the slice into. + index (int) : the index to select with. + name (str, optional): Name for the operation (optional, default is None). + + Returns: + Tensor, same dtype and shape with x + + Examples: + .. code-block:: python + + >>> import paddle + + >>> x = paddle.zeros((2,3,4)).astype("float32") + >>> values = paddle.ones((2,4)).astype("float32") + >>> res = paddle.select_scatter(x,values,1,1) + >>> print(res) + Tensor(shape=[2, 3, 4], dtype=float32, place=Place(cpu), stop_gradient=True, + [[[0., 0., 0., 0.], + [1., 1., 1., 1.], + [0., 0., 0., 0.]], + [[0., 0., 0., 0.], + [1., 1., 1., 1.], + [0., 0., 0., 0.]]]) + + """ + x_shape = x.shape + value_shape = values.shape + if not isinstance(x_shape, list): + x_shape = list(x_shape) + if index < 0: + index += x_shape[axis] + if axis < 0: + axis += len(x_shape) + del x_shape[axis] + if len(x_shape) != len(value_shape): + raise RuntimeError( + "expected values to have a size equal to the slice of x. value size = " + + str(value_shape) + + " slice size = " + + str(x_shape) + ) + for i in range(len(x_shape)): + if x_shape[i] != value_shape[i]: + raise RuntimeError( + "expected values to have a size equal to the slice of x. value size = " + + str(value_shape) + + " slice size = " + + str(x_shape) + ) + from ..base.framework import default_main_program + + starts = [index] + ends = [index + 1] + steps = [1] + axes = [axis] + none_axes = [] + decrease_axes = [axis] + inputs = {'Input': x} + attrs = { + 'axes': axes, + 'starts': starts, + 'ends': ends, + 'steps': steps, + 'decrease_axes': decrease_axes, + 'none_axes': none_axes, + } + + dtype = x.dtype + attrs['dtype'] = dtype + + values = values.astype(dtype) + inputs["ValueTensor"] = values + + if in_dynamic_or_pir_mode(): + return _C_ops.set_value_with_tensor( + x, + values, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + ) + else: + helper = LayerHelper('select_scatter', **locals()) + output = helper.create_variable_for_type_inference(dtype=x.dtype) + cur_block = default_main_program().current_block() + cur_block.append_op( + type="set_value", + inputs=inputs, + outputs={'Out': output}, + attrs=attrs, + inplace_map={"Input": "Out"}, + ) + + return output diff --git a/test/legacy_test/test_select_scatter_op.py b/test/legacy_test/test_select_scatter_op.py new file mode 100644 index 0000000000000..ecb5a0927e31b --- /dev/null +++ b/test/legacy_test/test_select_scatter_op.py @@ -0,0 +1,167 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import unittest + +import numpy as np + +import paddle +from paddle.framework import core +from paddle.pir_utils import test_with_pir_api + +paddle.enable_static() + + +class TestSelectScatterAPI(unittest.TestCase): + def setUp(self): + np.random.seed(0) + self.shape = [2, 3, 4] + self.type = np.float32 + self.x_np = np.random.random(self.shape).astype(self.type) + self.place = [paddle.CPUPlace()] + self.axis = 1 + self.index = 1 + self.value_shape = [2, 4] + self.value_np = np.random.random(self.value_shape).astype(self.type) + self.x_feed = copy.deepcopy(self.x_np) + if core.is_compiled_with_cuda(): + self.place.append(paddle.CUDAPlace(0)) + + def get_out_ref(self, out_ref, index, value_np): + for i in range(2): + for j in range(4): + out_ref[i, index, j] = value_np[i, j] + + @test_with_pir_api + def test_api_static(self): + paddle.enable_static() + + def run(place): + with paddle.static.program_guard(paddle.static.Program()): + x = paddle.static.data('Src', self.shape, self.type) + value = paddle.static.data( + 'Values', self.value_shape, self.type + ) + out = paddle.select_scatter(x, value, self.axis, self.index) + exe = paddle.static.Executor(place) + res = exe.run( + feed={ + 'Src': self.x_feed, + 'Values': self.value_np, + }, + fetch_list=[out], + ) + + out_ref = copy.deepcopy(self.x_np) + self.get_out_ref(out_ref, self.index, self.value_np) + for out in res: + np.testing.assert_allclose(out, out_ref, rtol=0.001) + + for place in self.place: + run(place) + + def test_api_dygraph(self): + def run(place): + paddle.disable_static(place) + x_tensor = paddle.to_tensor(self.x_np) + value_tensor = paddle.to_tensor(self.value_np) + out = paddle.select_scatter( + x_tensor, value_tensor, self.axis, self.index + ) + out_ref = copy.deepcopy(self.x_np) + self.get_out_ref(out_ref, self.index, self.value_np) + np.testing.assert_allclose(out.numpy(), out_ref, rtol=0.001) + + paddle.enable_static() + + for place in self.place: + run(place) + + +class TestSelectScatterAPICase2(TestSelectScatterAPI): + def setUp(self): + np.random.seed(0) + self.shape = [2, 3, 4, 5] + self.type = np.float64 + self.x_np = np.random.random(self.shape).astype(self.type) + self.place = [paddle.CPUPlace()] + self.axis = 2 + self.index = 1 + self.value_shape = [2, 3, 5] + self.value_np = np.random.random(self.value_shape).astype(self.type) + self.x_feed = copy.deepcopy(self.x_np) + if core.is_compiled_with_cuda(): + self.place.append(paddle.CUDAPlace(0)) + + def get_out_ref(self, out_ref, index, value_np): + for i in range(2): + for j in range(3): + for k in range(5): + out_ref[i, j, index, k] = value_np[i, j, k] + + +class TestSelectScatterAPICase3(TestSelectScatterAPI): + def setUp(self): + np.random.seed(0) + self.shape = [2, 3, 4, 5, 6] + self.type = np.int32 + self.x_np = np.random.random(self.shape).astype(self.type) + self.place = [paddle.CPUPlace()] + self.axis = 2 + self.index = 1 + self.value_shape = [2, 3, 5, 6] + self.value_np = np.random.random(self.value_shape).astype(self.type) + self.x_feed = copy.deepcopy(self.x_np) + if core.is_compiled_with_cuda(): + self.place.append(paddle.CUDAPlace(0)) + + def get_out_ref(self, out_ref, index, value_np): + for i in range(2): + for j in range(3): + for k in range(5): + for w in range(6): + out_ref[i, j, index, k, w] = value_np[i, j, k, w] + + +class TestSelectScatterAPIError(unittest.TestCase): + def setUp(self): + np.random.seed(0) + self.shape = [2, 3, 4] + self.x_np = np.random.random(self.shape).astype(np.float32) + self.place = [paddle.CPUPlace()] + self.axis = 1 + self.index = 1 + self.value_shape = [2, 4] + self.value_np = np.random.random(self.value_shape).astype(np.float32) + self.x_feed = copy.deepcopy(self.x_np) + if core.is_compiled_with_cuda(): + self.place.append(paddle.CUDAPlace(0)) + + def test_len_of_shape_not_equal_error(self): + with self.assertRaises(RuntimeError): + x_tensor = paddle.to_tensor(self.x_np) + value_tensor = paddle.to_tensor(self.value_np).reshape((2, 2, 2)) + res = paddle.select_scatter(x_tensor, value_tensor, 1, 1) + + def test_one_of_size_not_equal_error(self): + with self.assertRaises(RuntimeError): + x_tensor = paddle.to_tensor(self.x_np) + value_tensor = paddle.to_tensor([[2, 2], [2, 2]]).astype(np.float32) + res = paddle.select_scatter(x_tensor, value_tensor, 1, 1) + + +if __name__ == "__main__": + paddle.enable_static() + unittest.main()