Lightning-AI/lightning-thunder
Support `torch.where(condition)` with thunder.jit
オープン
#124 opened on 2024/04/03
enhancementhelp wantedoperators
Repository metrics
- Stars
- (1,460 個のスター)
- PR merge metrics
- (30d に merged PR はありません)
説明
🚀 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
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