Lightning-AI/lightning-thunder

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

開放

#124 建立於 2024年4月3日

在 GitHub 查看
 (12 則留言) (0 個反應) (0 位負責人)Python (114 個分叉)github user discovery
enhancementhelp wantedoperators

倉庫指標

星標
 (1,460 顆星)
PR 合併指標
 (30 天內沒有已合併 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 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

貢獻者指南

研究方向
檢查 thunder 程式碼庫中現有的 `torch.where(condition, input, other)` 實作，並擴展它以支援單參數變體 `torch.where(condition)`，該變體回傳非零元素的索引。
技術棧
pythonpytorch
領域
backend
議題類型
功能
難度面向新貢獻者的預計實作難度，1 表示很小改動，5 表示專家級工作。
3
預計時間有經驗貢獻者完成調查、實作、測試並準備 pull request 的粗略時間範圍。
半天
活動狀態議題目前的可參與程度：新鮮、活躍、陳舊、阻塞或等待維護者輸入。
活躍
清晰度議題是否清楚說明預期改動、驗收標準和下一步。
清晰
前置要求
PythonPyTorch
新手友善度1-100 的估計分數，表示該議題對首次貢獻者的友善程度。
60