Lightning-AI/lightning-thunder
[feature request] Indexing with boolean masks
Ouverte
#126 ouverte le 3 avr. 2024
enhancementhelp wantedoperators
Métriques du dépôt
- Stars
- (1 460 étoiles)
- Métriques de merge PR
- (Aucune PR mergée en 30 j)
Description
🚀 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