【Hackathon 5th No.26】为 Paddle 新增 diagonal_scatter API -part (#57879)

* add diagonal_scatter * add diagonal_scatter_test * add annotations * modify int32/int64 test range * update diagonal_scatter doc * fix: bfloat16 type & test name & LayerHelper position * add diagonal scatter in __all__ list * add int16 in fill_diagonal_tensor & reuse the logic of fill_diagonal_tensor for diagonal_scatter * remove xpu settings * op test get_numeric_gradient does not support int16 * fix * Update python/paddle/tensor/manipulation.py Co-authored-by: ooo oo <106524776+ooooo-create@users.noreply.github.com> * Update python/paddle/tensor/manipulation.py Co-authored-by: ooo oo <106524776+ooooo-create@users.noreply.github.com> * fix review suggestions * Update python/paddle/tensor/manipulation.py Co-authored-by: ooo oo <106524776+ooooo-create@users.noreply.github.com> * fix --------- Co-authored-by: ooo oo <106524776+ooooo-create@users.noreply.github.com>
PaddlePaddle · Nov 17, 2023 · 98bb0b8 · 98bb0b8
1 parent 00e5f28
commit 98bb0b8
Show file tree

Hide file tree

Showing 8 changed files with 479 additions and 1 deletion.
diff --git a/paddle/phi/kernels/cpu/fill_diagonal_tensor_grad_kernel.cc b/paddle/phi/kernels/cpu/fill_diagonal_tensor_grad_kernel.cc
@@ -76,6 +76,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor_grad,
                    double,
                    int64_t,
                    int,
+                   int16_t,
                    int8_t,
                    uint8_t,
                    phi::dtype::float16,

diff --git a/paddle/phi/kernels/cpu/fill_diagonal_tensor_kernel.cc b/paddle/phi/kernels/cpu/fill_diagonal_tensor_kernel.cc
@@ -140,6 +140,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor,
                    double,
                    int64_t,
                    int,
+                   int16_t,
                    int8_t,
                    uint8_t,
                    phi::dtype::float16,

diff --git a/paddle/phi/kernels/gpu/fill_diagonal_tensor_grad_kernel.cu b/paddle/phi/kernels/gpu/fill_diagonal_tensor_grad_kernel.cu
@@ -102,6 +102,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor_grad,
                    double,
                    int64_t,
                    int,
+                   int16_t,
                    int8_t,
                    uint8_t,
                    phi::dtype::float16,

diff --git a/paddle/phi/kernels/gpu/fill_diagonal_tensor_kernel.cu b/paddle/phi/kernels/gpu/fill_diagonal_tensor_kernel.cu
@@ -124,6 +124,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor,
                    double,
                    int64_t,
                    int,
+                   int16_t,
                    int8_t,
                    uint8_t,
                    phi::dtype::float16,

diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -260,6 +260,7 @@
     masked_fill_,
     index_fill,
     index_fill_,
+    diagonal_scatter,
 )
 
 from .tensor.math import (  # noqa: F401
@@ -928,4 +929,5 @@
     'hypot_',
     'index_fill',
     "index_fill_",
+    'diagonal_scatter',
 ]
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
@@ -173,6 +173,7 @@
 from .manipulation import masked_fill_  # noqa: F401
 from .manipulation import index_fill  # noqa: F401
 from .manipulation import index_fill_  # noqa: F401
+from .manipulation import diagonal_scatter  # noqa: F401
 from .math import abs  # noqa: F401
 from .math import abs_  # noqa: F401
 from .math import acos  # noqa: F401
@@ -741,6 +742,7 @@
     'atleast_1d',
     'atleast_2d',
     'atleast_3d',
+    'diagonal_scatter',
 ]
 
 # this list used in math_op_patch.py for magic_method bind

diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -1017,7 +1017,64 @@ def _fill_diagonal_tensor_impl(x, y, offset=0, dim1=0, dim2=1, inplace=False):
 
     if inplace:
         return _C_ops.fill_diagonal_tensor_(x, y, offset, dim1, dim2)
-    return _C_ops.fill_diagonal_tensor(x, y, offset, dim1, dim2)
+
+    if in_dynamic_mode():
+        return _C_ops.fill_diagonal_tensor(x, y, offset, dim1, dim2)
+    else:
+        check_variable_and_dtype(
+            x,
+            'X',
+            [
+                'float16',
+                'float32',
+                'float64',
+                'uint16',
+                'uint8',
+                'int8',
+                'int16',
+                'int32',
+                'int64',
+                'bool',
+                'complex64',
+                'complex128',
+            ],
+            'paddle.tensor.manipulation.fill_diagonal_tensor',
+        )
+        check_variable_and_dtype(
+            y,
+            'Y',
+            [
+                'float16',
+                'float32',
+                'float64',
+                'uint16',
+                'uint8',
+                'int8',
+                'int16',
+                'int32',
+                'int64',
+                'bool',
+                'complex64',
+                'complex128',
+            ],
+            'paddle.tensor.manipulation.fill_diagonal_tensor',
+        )
+        helper = LayerHelper('fill_diagonal_tensor', **locals())
+        out = helper.create_variable_for_type_inference(x.dtype)
+        helper.append_op(
+            type='fill_diagonal_tensor',
+            inputs={
+                'X': x,
+                'Y': y,
+            },
+            outputs={'Out': out},
+            attrs={
+                'offset': offset,
+                'dim1': dim1,
+                'dim2': dim2,
+            },
+        )
+        return out
 
 
 def fill_diagonal_tensor_(x, y, offset=0, dim1=0, dim2=1, name=None):
@@ -5770,3 +5827,46 @@ def index_fill_(x, index, axis, value, name=None):
 
     """
     return _index_fill_impl(x, index, axis, value, True)
+
+
+def diagonal_scatter(x, y, offset=0, axis1=0, axis2=1, name=None):
+    """
+    Embed the values of Tensor ``y`` into Tensor ``x`` along the diagonal elements
+    of Tensor ``x``, with respect to ``axis1`` and ``axis2``.
+
+    This function returns a tensor with fresh storage.
+
+    The argument ``offset`` controls which diagonal to consider:
+
+    - If ``offset`` = 0, it is the main diagonal.
+    - If ``offset`` > 0, it is above the main diagonal.
+    - If ``offset`` < 0, it is below the main diagonal.
+
+    Note:
+        ``y`` should have the same shape as :ref:`paddle.diagonal <api_paddle_diagonal>`.
+
+    Args:
+        x (Tensor): ``x`` is the original Tensor. Must be at least 2-dimensional.
+        y (Tensor): ``y`` is the Tensor to embed into ``x``
+        offset (int, optional): which diagonal to consider. Default: 0 (main diagonal).
+        axis1 (int, optional): first axis with respect to which to take diagonal. Default: 0.
+        axis2 (int, optional): second axis with respect to which to take diagonal. Default: 1.
+        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        Tensor, Tensor with diagonal embedeed with ``y``.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> x = paddle.arange(6.0).reshape((2, 3))
+            >>> y = paddle.ones((2,))
+            >>> out = x.diagonal_scatter(y)
+            >>> print(out)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(gpu:0), stop_gradient=True,
+                   [[1., 1., 2.],
+                    [3., 1., 5.]])
+
+    """
+    return fill_diagonal_tensor(x, y, offset, axis1, axis2, name)