Lightning-AI/lightning-thunder

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

Offen

#124 geöffnet am 03.04.2024

Auf GitHub ansehen
 (12 Kommentare) (0 Reaktionen) (0 zugewiesene Personen)Python (114 Forks)github user discovery
enhancementhelp wantedoperators

Repository-Metriken

Stars
 (1.460 Sterne)
PR-Merge-Metriken
 (Keine gemergten PRs in 30 T)

Beschreibung

🚀 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

Contributor Guide

Research-Richtung
Untersuchen Sie die vorhandene Implementierung von `torch.where(condition, input, other)` im Thunder Codebestand und erweitern Sie sie, um die einargumige Variante `torch.where(condition)` zu unterstützen, die Indizes von Nicht Null Elementen zurückgibt.
Tech Stack
pythonpytorch
Domain
backend
Issue Type
Funktion
SchwierigkeitGeschätzte Implementierungsschwierigkeit für neue Contributors, von 1 für sehr kleine Änderungen bis 5 für Expertenarbeit.
3
Geschätzte ZeitEin grober Zeitrahmen, um zu recherchieren, zu implementieren, zu testen und einen Pull Request vorzubereiten.
Halber Tag
AktivitätsstatusWie verfügbar das Issue gerade wirkt: frisch, aktiv, stale, blockiert oder wartet auf Maintainer-Input.
Aktiv
KlarheitWie klar das Issue die erwartete Änderung, Akzeptanzkriterien und den nächsten Schritt erklärt.
Klar
Voraussetzungen
PythonPyTorch
EinsteigerfreundlichkeitEin Score von 1 bis 100, der schätzt, wie zugänglich dieses Issue für First-Time Contributors ist.
60