Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 16 commits into from
Nov 17, 2023
Merged
1 change: 1 addition & 0 deletions paddle/phi/kernels/cpu/fill_diagonal_tensor_grad_kernel.cc
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor_grad,
double,
int64_t,
int,
int16_t,
int8_t,
uint8_t,
phi::dtype::float16,
Expand Down
1 change: 1 addition & 0 deletions paddle/phi/kernels/cpu/fill_diagonal_tensor_kernel.cc
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor,
double,
int64_t,
int,
int16_t,
int8_t,
uint8_t,
phi::dtype::float16,
Expand Down
1 change: 1 addition & 0 deletions paddle/phi/kernels/gpu/fill_diagonal_tensor_grad_kernel.cu
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor_grad,
double,
int64_t,
int,
int16_t,
int8_t,
uint8_t,
phi::dtype::float16,
Expand Down
1 change: 1 addition & 0 deletions paddle/phi/kernels/gpu/fill_diagonal_tensor_kernel.cu
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,7 @@ PD_REGISTER_KERNEL(fill_diagonal_tensor,
double,
int64_t,
int,
int16_t,
int8_t,
uint8_t,
phi::dtype::float16,
Expand Down
2 changes: 2 additions & 0 deletions python/paddle/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,7 @@
masked_fill_,
index_fill,
index_fill_,
diagonal_scatter,
DanGuge marked this conversation as resolved.
Show resolved Hide resolved
)

from .tensor.math import ( # noqa: F401
Expand Down Expand Up @@ -924,4 +925,5 @@
'hypot_',
'index_fill',
"index_fill_",
'diagonal_scatter',
]
2 changes: 2 additions & 0 deletions python/paddle/tensor/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -737,6 +738,7 @@
'atleast_1d',
'atleast_2d',
'atleast_3d',
'diagonal_scatter',
]

# this list used in math_op_patch.py for magic_method bind
Expand Down
102 changes: 101 additions & 1 deletion python/paddle/tensor/manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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:
DanGuge marked this conversation as resolved.
Show resolved Hide resolved

- 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)
Loading