【Hackathon 5th No.26】为 Paddle 新增 diagonal_scatter API

[DanGuge]

为 Paddle 新增 diagonal_scatter API

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请检查PR提交格式和内容是否完备，具体请参考示例和模版。
Your PR has been submitted. Thanks for your contribution!
Please check its format and content. For this, you can refer to Template and Demo.

[CLAassistant]

All committers have signed the CLA.

[DanGuge]

补充一下MindSpore上的实现，该目前应该只支持方阵，受限于下图红框中的运算：

link：https://github.com/mindspore-ai/mindspore/blob/master/mindspore/python/mindspore/ops/function/array_func.py#L6942

该方法试图保留input矩阵的对角线上的元素，但是ones_like(src)和input的形状使该运算失败

测试

• 输入矩阵x的shape为(2, 3)，矩阵y的shape(2, )

• 在mindspore 2.1.1上的测试
  

• 在torch 2.0.1上的测试
  

[zxcd]

inplace写是指？

[DanGuge]

• 实现的逻辑是这样的，我们通过diagonal()函数获得output矩阵（x矩阵的clone）的diagonal_slice视图
  • 在动态图中，我们可以直接通过张量索引将y矩阵的元素写入diagonal_slice视图，这样能够直接让output矩阵中对应元素也发生修改
  • 在静态图中，设置元素需要通过paddle.static.setitem()函数实现，但是它在静态图中的逻辑不是直接修改diagonal_slice视图中的元素，而是返回一个新的矩阵，这样就没法实现对output矩阵的修改
• paddle.static.setitem(x, index, value)的调用最终会调用_setitem_static(x, indices, values)函数
• 函数链接如下：
  • https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/static/input.py#L401
  • https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/base/variable_index.py#L809

[zxcd]

如果使用setitem返回一个新的矩阵，再通过Variable中的set_value函数向原矩阵赋值的做法是否能够满足需求？

[DanGuge]

def diagonal_scatter(x, y, offset=0, axis1=0, axis2=1, name=None):
    output = x.clone()
    diagonal_slice = output.diagonal(offset, axis1, axis2)
    if diagonal_slice.shape != y.shape:
        raise ValueError(
            "The shape of diagonal_slice is not equal to the shape of y. diagonal_slice shape = {}, y shape = {}".format(
                diagonal_slice.shape, y.shape
            )
        )
    if in_dynamic_mode():
        diagonal_slice[:] = y
    else:
        from builtins import slice
        diagonal_slice = paddle.static.setitem(diagonal_slice, slice(None, None, None), y)
    return output

我们的理解上应该有些误差，这个是我目前的实现方案：

• 在动态图中，结果符合预期
  • diagonal_slice是x矩阵diagonal的视图，使用diagonal_slice[:] = y，可以直接对视图中元素进行修改，这样x矩阵的值也会改变
• 在静态图中，我也想做相同的事情，就是拿到x矩阵的diagonal视图，通过对视图元素的修改，引起源矩阵x的修改
  • 但是paddle.static.setitem函数在这里不会直接对diagonal_slice视图中的元素修改，而是返回一个新的diagonal_slice矩阵

所以并不是新的矩阵符合需求，而是没法通过这种方式得到符合需求的x矩阵

[zxcd]

我感觉这样做可能会有问题，你尝试过这样修改在动态图是可行的吗？

[DanGuge]

在动态图上测试是可行的，这个是目前的PR：PaddlePaddle/Paddle#57879
在torch的实现中，整体逻辑也是类似的：

• cpp实现，https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/TensorShape.cpp#L3971

at::Tensor diagonal_scatter(const at::Tensor& self, const at::Tensor& src, int64_t offset, int64_t dim1, int64_t dim2) {
    // See Note [*_scatter ops preserve strides]
    auto output = clone_preserve_strides(self);
    auto slice = output.diagonal(offset, dim1, dim2);
    TORCH_CHECK(slice.sizes() == src.sizes(), "expected src to have a size equal to the slice of self. src size = ", src.sizes(), ", slice size = ", slice.sizes());
    slice.copy_(src);
    return output;
}

• python新后端的逻辑，https://github.com/pytorch/pytorch/blob/main/torch/_inductor/lowering.py#L1023

@register_lowering(aten.diagonal_scatter, type_promotion_kind=None)
def diagonal_scatter(input, src, offset: int = 0, dim1: int = 0, dim2: int = 1):
    output = clone(input)
    target = diagonal(output, offset, dim1, dim2)
    mutate_to(target, src)
    return output

[zxcd]

静态图是不支持view，所以这样确实不行。另外看了一下pytorch，diagal_scatter是返回一个新的tensor，并不是在原值上修改，设计的时候返回一个新的tensor就可以了。这样动静态图都可以支持。

[DanGuge]

此处使用view的方式来修改元素，是因为没有python中没有特别好的能拿到diagonal的index的方式，所以只好通过view的方式来修改，torch中也是通过调用diagonal函数拿到slice，然后直接将src矩阵的元素写入slice中

或者我可以通过复用_C_ops.fill_diagonal_tensor(x, y, offset, dim1, dim2)的实现，来实现diagonal_scatter吗？我看在 https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/tensor/manipulation.py#L1019 描述的基本操作逻辑和diagonal_scatter是一致的

[zxcd]

看逻辑应该可以试试，辛苦修改一下方案

[zxcd]

fill_diagonal_tensor也支持复数类型，这个类型是否也能支持？

[DanGuge]

应该也是支持的，已修改设计文档

[zxcd]

以下是fill_diagonal_tensor GPU 和CPU支持的数据类型，烦请参考这里将类型补全吧。

PD_REGISTER_KERNEL(fill_diagonal_tensor,
                   GPU,
                   ALL_LAYOUT,
                   phi::FillDiagonalTensorKernel,
                   float,
                   double,
                   int64_t,
                   int,
                   int8_t,
                   uint8_t,
                   phi::dtype::float16,
                   phi::dtype::bfloat16,
                   phi::dtype::complex<float>,
                   phi::dtype::complex<double>,
                   bool) {}

PD_REGISTER_KERNEL(fill_diagonal_tensor,
                   CPU,
                   ALL_LAYOUT,
                   phi::FillDiagonalTensorKernel,
                   float,
                   double,
                   int64_t,
                   int,
                   int8_t,
                   uint8_t,
                   phi::dtype::float16,
                   phi::dtype::complex<float>,
                   phi::dtype::complex<double>,
                   bool) {}

[DanGuge]

@zxcd dtype已经补全完了，辛苦review一下

[zxcd]

LGTM