fix behavior of put_along_axis and take_along_axis 易用性提升No.43

[YibinLiu666]

PR types

Others

PR changes

Others

Description

#55883
该PR修改内容如下：

1. 【功能增强】：当前是index必须broadcast的版本，在此基础上额外支持了一种非broadcast版本的scatter，以适应更多场景的灵活使用。
2. 【Bug修复】：修复了gpu上reduce到相同位置时与np.put_along_axis以及torch.scatter结果都不一致的bug。
3. 【Bug修复】：修复了reduce为'assign'时，scatter到相同位置时value梯度计算错误的bug。

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[paddle-bot]

✅ This PR's description meets the template requirements!
Please wait for other CI results.

[YibinLiu666]

PR-CE-Framework 没过应该是paddletest里面的put_along_axis以及take_along_axis的测试样例仍然是与numpy对齐导致过不了。但是修改后的算子相比之前在cpu上的性能下降了，我自己能想到优化的地方都进行了优化，能否麻烦reviewer帮忙看看还有哪里可以再改进的。

[zhwesky2010]

@YibinLiu666 这里pytorch使用的是 index不向输入broadcast 的方式，numpy和paddle使用的是 index向输入broadcast 的方式，因此如果 index和输入没有shape对齐 的情况下，两者结果不一致。

这里就有两个实现版本：1）broadcast版本；2）非broadcast版本。

目前的改动是 直接改成pytorch的非broadcast版本，会有很大的不兼容影响，我建议是否可以增加一个参数：broadcast=True，默认为True时仍是原本的broadcast版本逻辑，以保证兼容；如果设置broadcast=False，则新增实现 pytorch的非broadcast版本。

[YibinLiu666]

@YibinLiu666 这里pytorch使用的是 index不向输入broadcast 的方式，numpy和paddle使用的是 index向输入broadcast 的方式，因此如果 index和输入没有shape对齐 的情况下，两者结果不一致。

这里就有两个实现版本：1）broadcast版本；2）非broadcast版本。

目前的改动是 直接改成pytorch的非broadcast版本，会有很大的不兼容影响，我建议是否可以增加一个参数：broadcast=True，默认为True时仍是原本的broadcast版本逻辑，以保证兼容；如果设置broadcast=False，则新增实现 pytorch的非broadcast版本。

已经修改，改动后的kernel是跟以前兼容的，因此只在顶层python接口处添加了broadcast参数

[zhwesky2010]

目前像这样改了后，还会影响性能吗，理论上不应该影响到性能。

[YibinLiu666]

目前像这样改了后，还会影响性能吗，理论上不应影响性能

目前CI中没有明显的性能下降，只不过应该还是会有一点点影响，因为之前是默认value的shape跟index的是一样的，而现在不一定一样的，在kernel里面每次还额外计算了value的offset。然后GPU实现里面加了同步操作，这个是因为有时候会对src的同一个位置进行多次修改，比如输入的src是[[0,0,0],[0,0,0]]，index是[[0,0,0],[0,0,0]，value的值是[[1,2,3],[4,5,6]]，axis=1的时候，当i都是0时，index[0,0]和index[0,1]的值都是0，因此src要修改的位置都会是src[0,0]，如果没有同步操作的话可能会存在随机性，因为不知道哪个gpu线程先结束，pytorch的行为是保留i,j字典序最大的那个操作，所以我这里同步了一下，只让tid最大的线程修改src

[luotao1]

目前CI中没有明显的性能下降

Op-Benchmark流水线能通过说明没有5%的性能下降，我看某些case还快了

[zhwesky2010]

目前像这样改了后，还会影响性能吗，理论上不应影响性能

目前CI中没有明显的性能下降，只不过应该还是会有一点点影响，因为之前是默认value的shape跟index的是一样的，而现在不一定一样的，在kernel里面每次还额外计算了value的offset。然后GPU实现里面加了同步操作，这个是因为有时候会对src的同一个位置进行多次修改，比如输入的src是[[0,0,0],[0,0,0]]，index是[[0,0,0],[0,0,0]，value的值是[[1,2,3],[4,5,6]]，axis=1的时候，当i都是0时，index[0,0]和index[0,1]的值都是0，因此src要修改的位置都会是src[0,0]，如果没有同步操作的话可能会存在随机性，因为不知道哪个gpu线程先结束，pytorch的行为是保留i,j字典序最大的那个操作，所以我这里同步了一下，只让tid最大的线程修改src

如果broadcast了，那应该与原本逻辑一致，不应有任何性能上的改变呢。如果没有broadcast，才是新的逻辑。

[YibinLiu666]

目前像这样改了后，还会影响性能吗，理论上不应影响性能

目前CI中没有明显的性能下降，只不过应该还是会有一点点影响，因为之前是默认value的shape跟index的是一样的，而现在不一定一样的，在kernel里面每次还额外计算了value的offset。然后GPU实现里面加了同步操作，这个是因为有时候会对src的同一个位置进行多次修改，比如输入的src是[[0,0,0],[0,0,0]]，index是[[0,0,0],[0,0,0]，value的值是[[1,2,3],[4,5,6]]，axis=1的时候，当i都是0时，index[0,0]和index[0,1]的值都是0，因此src要修改的位置都会是src[0,0]，如果没有同步操作的话可能会存在随机性，因为不知道哪个gpu线程先结束，pytorch的行为是保留i,j字典序最大的那个操作，所以我这里同步了一下，只让tid最大的线程修改src

如果broadcast了，那应该与原本逻辑一致，不应有任何性能上的改变呢。如果没有broadcast，才是新的逻辑。

原本的逻辑我认为是存在问题的，首先就是我这里提到的gpu kernel，其次就是求value梯度的过程，之前就是用的gather来求梯度，但是还是相同的情况，如果多个index同时修改src的同一个地址，前面修改的会被最后的修改覆盖，也就是不会用到相关的value，梯度就为0，但是用gather的话被覆盖的value也会有梯度

[luotao1]

@YibinLiu666 有兴趣把黑客松第6题 【为Paddle增强 put_along_axis API】一起做了么

[YibinLiu666]

目前像这样改了后，还会影响性能吗，理论上不应影响性能

目前CI中没有明显的性能下降，只不过应该还是会有一点点影响，因为之前是默认value的shape跟index的是一样的，而现在不一定一样的，在kernel里面每次还额外计算了value的offset。然后GPU实现里面加了同步操作，这个是因为有时候会对src的同一个位置进行多次修改，比如输入的src是[[0,0,0],[0,0,0]]，index是[[0,0,0],[0,0,0]，value的值是[[1,2,3],[4,5,6]]，axis=1的时候，当i都是0时，index[0,0]和index[0,1]的值都是0，因此src要修改的位置都会是src[0,0]，如果没有同步操作的话可能会存在随机性，因为不知道哪个gpu线程先结束，pytorch的行为是保留i,j字典序最大的那个操作，所以我这里同步了一下，只让tid最大的线程修改src

如果broadcast了，那应该与原本逻辑一致，不应有任何性能上的改变呢。如果没有broadcast，才是新的逻辑。

原本的逻辑我认为是存在问题的，首先就是我这里提到的gpu kernel，其次就是求value梯度的过程，之前就是用的gather来求梯度，但是还是相同的情况，如果多个index同时修改src的同一个地址，前面修改的会被最后的修改覆盖，也就是不会用到相关的value，梯度就为0，但是用gather的话被覆盖的value也会有梯度

比如这个样例，pytorch的输出是[[11,12,13,14,15],[0,0,0,0,0],[0,0,0,0,0]]，此外，这里只用到了1,2,3,4,5，但是value的梯度全是1

[YibinLiu666]

@YibinLiu666 有兴趣把黑客松第6题 【为Paddle增强 put_along_axis API】一起做了么

可以呀，只不过我现在还在写27题

[zhwesky2010]

目前像这样改了后，还会影响性能吗，理论上不应影响性能

目前CI中没有明显的性能下降，只不过应该还是会有一点点影响，因为之前是默认value的shape跟index的是一样的，而现在不一定一样的，在kernel里面每次还额外计算了value的offset。然后GPU实现里面加了同步操作，这个是因为有时候会对src的同一个位置进行多次修改，比如输入的src是[[0,0,0],[0,0,0]]，index是[[0,0,0],[0,0,0]，value的值是[[1,2,3],[4,5,6]]，axis=1的时候，当i都是0时，index[0,0]和index[0,1]的值都是0，因此src要修改的位置都会是src[0,0]，如果没有同步操作的话可能会存在随机性，因为不知道哪个gpu线程先结束，pytorch的行为是保留i,j字典序最大的那个操作，所以我这里同步了一下，只让tid最大的线程修改src

如果broadcast了，那应该与原本逻辑一致，不应有任何性能上的改变呢。如果没有broadcast，才是新的逻辑。

原本的逻辑我认为是存在问题的，首先就是我这里提到的gpu kernel，其次就是求value梯度的过程，之前就是用的gather来求梯度，但是还是相同的情况，如果多个index同时修改src的同一个地址，前面修改的会被最后的修改覆盖，也就是不会用到相关的value，梯度就为0，但是用gather的话被覆盖的value也会有梯度

比如这个样例，pytorch的输出是[[11,12,13,14,15],[0,0,0,0,0],[0,0,0,0,0]]，此外，这里只用到了1,2,3,4,5，但是value的梯度全是1

感谢分析，这么看这个API确实问题比较多，上面这个case np.put_along_axis的结果和torch也是一致的，而paddle是错误的。现在有三个方面的修复或增强：

1. put_along_axis在index重复，也就是存在reduce的情况下，前向值不对。反向值应该根据reduce的设置来定，如果是"assign"，则后面的梯度为1，如果是"add"，则梯度均为1，如果是"multiply"，则梯度应该为部分元素的乘积；
2. put_along_axis目前采用broadcast的方式，对于 非broadcast的功能需要增强，也就是增强torch.scatter的功能，增加一个broadcast参数；
3. 可见黑客松第6题 【为Paddle增强 put_along_axis API】，需要增强torch.scatter_reduce的功能，也就是reduce的支持情况，目前支持的reduce类型相比torch.scatter_reduce少3种，另外缺少参数include_self的参数功能。

为了避免后面的不兼容，建议三个功能整体考虑，最终需要增加两个参数：include_self=True, broadcast=True，如果这个PR先合入的话，建议在broadcast=True前面补上include_self=True作为占位符，include_self=False目前没有实现先抛出异常。

[YibinLiu666]

目前像这样改了后，还会影响性能吗，理论上不应影响性能

目前CI中没有明显的性能下降，只不过应该还是会有一点点影响，因为之前是默认value的shape跟index的是一样的，而现在不一定一样的，在kernel里面每次还额外计算了value的offset。然后GPU实现里面加了同步操作，这个是因为有时候会对src的同一个位置进行多次修改，比如输入的src是[[0,0,0],[0,0,0]]，index是[[0,0,0],[0,0,0]，value的值是[[1,2,3],[4,5,6]]，axis=1的时候，当i都是0时，index[0,0]和index[0,1]的值都是0，因此src要修改的位置都会是src[0,0]，如果没有同步操作的话可能会存在随机性，因为不知道哪个gpu线程先结束，pytorch的行为是保留i,j字典序最大的那个操作，所以我这里同步了一下，只让tid最大的线程修改src

如果broadcast了，那应该与原本逻辑一致，不应有任何性能上的改变呢。如果没有broadcast，才是新的逻辑。

原本的逻辑我认为是存在问题的，首先就是我这里提到的gpu kernel，其次就是求value梯度的过程，之前就是用的gather来求梯度，但是还是相同的情况，如果多个index同时修改src的同一个地址，前面修改的会被最后的修改覆盖，也就是不会用到相关的value，梯度就为0，但是用gather的话被覆盖的value也会有梯度

比如这个样例，pytorch的输出是[[11,12,13,14,15],[0,0,0,0,0],[0,0,0,0,0]]，此外，这里只用到了1,2,3,4,5，但是value的梯度全是1

感谢分析，这么看这个API确实问题比较多，上面这个case np.put_along_axis的结果和torch也是一致的，而paddle是错误的。现在有三个方面的修复或增强：

1. put_along_axis在index重复，也就是存在reduce的情况下，前向值不对。反向值应该根据reduce的设置来定，如果是"assign"，则后面的梯度为1，如果是"add"，则梯度均为1，如果是"multiply"，则梯度应该为部分元素的乘积；
2. put_along_axis目前采用broadcast的方式，对于 非broadcast的功能需要增强，也就是增强torch.scatter的功能，增加一个broadcast参数；
3. 可见黑客松第6题 【为Paddle增强 put_along_axis API】，需要增强torch.scatter_reduce的功能，也就是reduce的支持情况，目前支持的reduce类型相比torch.scatter_reduce少3种，另外缺少参数include_self的参数功能。

为了避免后面的不兼容，建议三个功能整体考虑，最终需要增加两个参数：include_self=True, broadcast=True，如果这个PR先合入的话，建议在broadcast=True前面补上include_self=True作为占位符，include_self=False目前没有实现先抛出异常。

好的，我再改改，梯度bug和功能增强我在黑客松第6题一起解决

[YibinLiu666]

@zhwesky2010 这个PR能麻烦尽早merge嘛，后面黑客松的第6题得基于这个PR，感谢

[luotao1]

有检查异常的单测么，覆盖率不够

[YibinLiu666]

有检查异常的单测么，覆盖率不够

已添加

[zhwesky2010]

LGTM

[YibinLiu666]

@luotao1 涛姐这个PR能先merge了吗，周末不小心把其他代码push到这个分支了，已经revert到研发大哥approve的那个版本了

[luotao1]

@YibinLiu666 请提交对应的中文文档修改PR

[zhwesky2010]

@YibinLiu666 你好，这里报出了一些不兼容问题。麻烦看下。

[zhwesky2010]

@YibinLiu666 另外看这里的中文文档也是有问题的 https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/put_along_axis_cn.html#put-along-axis

[YibinLiu666]

@YibinLiu666 你好，这里报出了一些不兼容问题。麻烦看下。

能否提供一下这两个pdt文件

[YibinLiu666]

@YibinLiu666 你好，这里报出了一些不兼容问题。麻烦看下。

目前发现是GPU上共享内存最大只能分配100KB左右导致的，数据量太大导致共享内存不够用。我明天提个PR修复一下。

[zhwesky2010]

@YibinLiu666 API文档也都修一下吧

[zhwesky2010]

@YibinLiu666 你好，上面的计算问题已经验证修复了。但是发现当shape很大时，性能下降非常多，不是之前说的只有一点点影响：

这个是测试方式，可以看出，反向没有问题，但put_along_axis的3个PR合入后前向的性能变慢了400-500倍，辛苦尽快看一下这个问题，影响比较大~

[SigureMo]

在静态图模式下，indices 是一个 Variable/Value，组网阶段是没有值的，这里 not Value() 永远返回 False，也就是永远不会走到下面的分支

我们正在完全禁掉 Value.__bool__ 以免发生一些低级错误，目前会在本 PR 改动的这里报错，麻烦修改一下相关代码，不要触发隐式 Value.__bool__

详情见 #60902，可拉取 PR 后运行 test/legacy_test/test_put_along_axis_op.py 复现

考虑到这里组网阶段是无法得知结果的，这个检查可能只能在动态图做，可以考虑改成

-if not (indices < axis_max_size).all():
+if in_dynamic_mode() and not (indices < axis_max_size).all():