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

貢獻者指南