Lightning-AI/lightning-thunder

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

Aberta

#124 aberto em 3 de abr. de 2024

Ver no GitHub
 (12 comentários) (0 reação) (0 responsável)Python (114 forks)github user discovery
enhancementhelp wantedoperators

Métricas do repositório

Stars
 (1.460 estrelas)
Métricas de merge de PR
 (Nenhuma PRs mesclada em 30d)

Description

🚀 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

Guia do colaborador

Direção de pesquisa
Examine a implementação existente de `torch.where(condition, input, other)` na base de código do thunder e estenda a para suportar a variante de um argumento `torch.where(condition)`, que retorna os índices dos elementos não nulos.
Pilha de tecnologia
pythonpytorch
Domain
backend
Tipo Issue
Funcionalidade
DifficultyDificuldade de implementação estimada para um novo contribuidor, de 1 para alterações muito pequenas a 5 para trabalho de nível especializado.
3
Tempo estimadoUm intervalo de tempo aproximado para um colaborador experiente investigar, implementar, testar e preparar um pull request.
Meio dia
Status da atividadeQuão disponível o issue aparece agora: novo, ativo, obsoleto, bloqueado ou aguardando entrada do mantenedor.
Ativo
ClarityCom que clareza o issue explica a mudança esperada, os critérios de aceitação e a próxima etapa.
Claro
Prerequisites
PythonPyTorch
Simpatia para novatosUma pontuação de 1 a 100 que estima o quão acessível este issue é para colaboradores iniciantes.
60