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

你的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.

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

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

@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 这里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参数

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

目前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

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

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

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

目前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，才是新的逻辑。

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

目前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也会有梯度

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

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

目前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也会有梯度

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

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

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

目前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目前没有实现先抛出异常。

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

目前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题一起解决

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

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

已添加

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

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

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

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

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

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

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

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

@YibinLiu666 API文档也都修一下吧

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

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