【Hackathon 5th No.26】Add diagonal_scatter API to Paddle

python/paddle/__init__.py

@@ -211,6 +211,7 @@
     scatter_,
     scatter_nd_add,
     scatter_nd,
+    diagonal_scatter,
     shard_index,
     slice,
     crop,
@@ -563,6 +564,7 @@
     'diagflat',
     'isnan',
     'scatter_nd_add',
+    'diagonal_scatter',
     'unstack',
     'get_default_dtype',
     'save',

python/paddle/tensor/__init__.py

@@ -130,6 +130,7 @@
 from .manipulation import scatter_  # noqa: F401
 from .manipulation import scatter_nd_add  # noqa: F401
 from .manipulation import scatter_nd  # noqa: F401
+from .manipulation import diagonal_scatter  # noqa: F401
 from .manipulation import shard_index  # noqa: F401
 from .manipulation import slice  # noqa: F401
 from .manipulation import split  # noqa: F401
@@ -380,6 +381,15 @@
 from ..signal import istft  # noqa: F401
 from ..signal import stft  # noqa: F401
 
+from ..signal import istft  # noqa: F401
+from ..signal import stft  # noqa: F401
+
+from ..signal import istft  # noqa: F401
+from ..signal import stft  # noqa: F401
+
+from ..signal import istft  # noqa: F401
+from ..signal import stft  # noqa: F401
+
 # this list used in math_op_patch.py for _binary_creator_
 tensor_method_func = [  # noqa
     'create_parameter',
@@ -558,6 +568,7 @@
     'scatter',
     'scatter_',
     'scatter_nd_add',
+    'diagonal_scatter',
     'scatter_nd',
     'shard_index',
     'slice',

python/paddle/tensor/manipulation.py

@@ -1975,8 +1975,10 @@ def split(x, num_or_sections, axis=0, name=None):
             )
 
         if isinstance(num_or_sections, int):
+            dim = dim if dim >= 0 else dim + len(input.shape)
             return _C_ops.split_with_num(input, num_or_sections, dim)
         else:
+            dim = dim if dim >= 0 else dim + len(input.shape)
             return _C_ops.split(input, num_or_sections, dim)
     elif in_pir_mode():
         if isinstance(dim, int):
@@ -5154,6 +5156,27 @@ def unfold(x, axis, size, step, name=None):
     """
     return _C_ops.tensor_unfold(x, axis, size, step)
 
+def _check_diagonal_scatter_shape(diag_shape, src_shape):
+    if diag_shape != src_shape:
+        raise ValueError(f"For diagonal_scatter, the shape of src should equal to the shape of input diagonal,"
+                         f"but got src.shape {src_shape} and diagonal shape {diag_shape}.")
+
+def diagonal_scatter(x, src, offset=0, dim1=0, dim2=1):
+    op_type = 'diagonal_scatter'
+    # input_dtype = ['int32', 'int64', 'float32', 'float64']
+
+    # check_variable_and_dtype(x, 'x', input_dtype, op_type)
+    # check_variable_and_dtype(src, 'src', input_dtype, op_type)
+    check_type(offset, 'offset', (int), op_type)
+    check_type(dim1, 'dim1', (int), op_type)
+    check_type(dim2, 'dim2', (int), op_type)
+    input_diag = paddle.diagonal(x, offset, dim1, dim2)
+    _check_diagonal_scatter_shape(input_diag.shape, src.shape)
+    embed = paddle.ones_like(src)
+    embed = paddle.nn.functional.diag_embed(embed, offset, dim1, dim2)
+    embed = x * embed
+    src = paddle.nn.functional.diag_embed(src, offset, dim1, dim2)
+    return x + src - embed
 
 # TODO(dev): We need avoid implementing it by this way.
 __METHODS = {

test/legacy_test/test_diagonal_scatter.py

@@ -0,0 +1,73 @@
+#  Copyright (c) 2023 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 unittest
+import numpy as np
+import paddle
+
+class TestDiagnalScatterAPI(unittest.TestCase):
+    def setUp(self):
+
+        self.x_dtype = 'float32'
+        self.x_shape = [3, 3]
+        self.x = np.zeros([3, 3]).astype('float32')
+        self.src_dtype = 'float32'
+        self.src_shape = [3]
+        self.src  = np.ones(3).astype('float32')
+        self.offset  = 0
+        self.init_input()
+        self.place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+    def init_input(self):
+        pass
+
+    def test_static_api(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(paddle.static.Program()):
+            x = paddle.static.data('x', self.x_shape, dtype=self.x_dtype)
+            src = paddle.static.data('src', self.src_shape, dtype=self.x_dtype)
+            out = paddle.diagonal_scatter(x, src, self.offset)
+            exe = paddle.static.Executor(self.place)
+            res = exe.run(
+                feed={
+                    'x': self.x,
+                    'src': self.src,
+                },
+                fetch_list=[out],
+            )
+            expect_output = [[1.0,0.0,0.0],
+                [0.0,1.0,0.0],
+                [0.0,0.0,1.0]]
+            np.testing.assert_allclose(expect_output, res[0], atol=1e-07)
+
+    def test_dygraph_api(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.x)
+        src = paddle.to_tensor(self.src)
+        res = paddle.diagonal_scatter(x, src, self.offset)
+
+        expect_output = np.array([[1.0,0.0,0.0],
+                [0.0,1.0,0.0],
+                [0.0,0.0,1.0]])
+        np.testing.assert_allclose(expect_output, res, atol=1e-07)
+        paddle.enable_static()
+
+
+if __name__ == '__main__':
+    paddle.enable_static()
+    unittest.main()