希望能添加广播索引机制

[BAYEKHELIX]

在抄informer的官方代码的时候，发现torch的广播索引机制在jittor中并不支持。 希望后续版本能添加一下，让抄代码更方便。

torch的原始代码如下

Q_reduce = Q[ torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], M_top, :] # factor*ln(L_q)

三个高维索引的形状分别是(B,1,1)、(1,H,1)、(B,H,X)，torch的话会进行广播把三个高维索引的维度广播到(B,H,X)进行索引。 目前jittor似乎没有这个机制，或者说有错误。

测试用例如下

然后是错误信息

目前采用手工广播实现目标功能，但我本身对torch了解不多，所以也不知道这样是不是torch的本意，只是能让代码跑起来了。

[Jittor]

收到，谢谢您的反馈😁 我们尽快支持

---原始邮件--- 发件人: @.> 发送时间: 2022年3月25日(周五) 中午11:37 收件人: @.>; 抄送: @.***>; 主题: [Jittor/jittor] 希望能添加广播索引机制 (Issue #301)

在抄informer的官方代码的时候，发现torch的广播索引机制在jittor中并不支持。 希望后续版本能添加一下，让抄代码更方便。

torch的原始代码如下

Q_reduce = Q[ torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], M_top, :] # factor*ln(L_q)

三个高维索引的形状分别是(B,1,1)、(1,H,1)、(B,H,X)，torch的话会进行广播把三个高维索引的维度广播到(B,H,X)进行索引。 目前jittor似乎没有这个机制，或者说有错误。

测试用例如下

然后是错误信息

目前采用手工广播实现目标功能，但我本身对torch了解不多，所以也不知道这样是不是torch的本意，只是能让代码跑起来了。

— Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you are subscribed to this thread.Message ID: @.***>

[cjld]

您好，jittor已经在最新版本1.3.3.8中添加了以上支持，相关commit：https://github.com/Jittor/jittor/commit/0b13930ed3b6c6d7ccc0e80adf610e4c28866a7f