Lightning-AI/lightning-thunder

Support `torch.where(condition)` with thunder.jit

Aperta

#124 aperta il 3 apr 2024

Vedi su GitHub
 (12 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

Motivation

Mixtral uses it: https://github.com/mistralai/mistral-src/blob/8598cf582091a596671be31990448e0620017851/moe_one_file_ref.py#L215

Minimal Repro

import thunder

def fn(cond):
    return torch.where(cond)

thunder.jit(fn)(torch.randn(3) > 0)

Pitch

Support image from https://pytorch.org/docs/stable/generated/torch.where.html

Additional context

We already support torch.where(condition, input, other): https://github.com/search?q=repo%3ALightning-AI%2Flightning-thunder+%22def+where%22&type=code

cc @apaz-cli

Guida contributor

Direzione di ricerca
Esamina l'implementazione esistente di `torch.where(condition, input, other)` nel codice di thunder ed estendila per supportare la variante a un argomento `torch.where(condition)`, che restituisce gli indici degli elementi diversi da zero.
Tech stack
pythonpytorch
Dominio
backend
Tipo issue
Funzionalità
DifficoltàQuanto dovrebbe essere impegnativa per un nuovo contributor.
3
Tempo stimatoTempo stimato per completare e verificare uno scope piccolo.
Mezza giornata
Stato attivitàQuanto è probabile che la issue sia ancora attiva e reviewable.
Attiva
ChiarezzaQuanto chiaramente la issue descrive problema, scope e risultato atteso.
Chiara
Prerequisiti
PythonPyTorch
Adatta ai principiantiQuanto la issue è adatta a contributor alla prima esperienza.
60