Lightning-AI/lightning-thunder

[feature request] Indexing with boolean masks

Aperta

#126 aperta il 3 apr 2024

 (0 commenti) (0 reazioni) (0 assegnatari)Python (114 fork)github user discovery
enhancementhelp wantedoperators

Metriche repository

Star
 (1460 stelle)
Metriche merge PR
 (Nessuna PR mergiata in 30 g)

Descrizione

🚀 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

Guida contributor