Lightning-AI/lightning-thunder

[feature request] Indexing with boolean masks

开放

#126 创建于 2024年4月3日

 (0 条评论) (0 个反应) (0 位负责人)Python (114 个派生)github user discovery
enhancementhelp wantedoperators

仓库指标

星标
 (1,460 个星标)
PR 合并指标
 (30 天内没有已合并 PR)

描述

🚀 Feature

No indexing with boolean mask, for example:

import torch; import thunder                                                                                   
m = x <= 0.5
                                                                                                                                                                                                                          
def f(x, m):                                                                                                   
    return x[m]                                    
                                                                                                        
                                                                                                                       
jf = thunder.jit(f)                                
jf(x, m)  

fails with

RuntimeError: Advanced indexing currently only supports zero or one-dimensional integer tensors, but found a tensor with dtype bool8 and 1 dimensions

cc @apaz-cli

贡献者指南