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

仓库指标

星标

1,460 个星标 (1,460 个星标)

PR 合并指标

30 天内没有已合并 PR (30 天内没有已合并 PR)

描述

import thunder

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

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

贡献者指南

研究方向

检查 thunder 代码库中现有的 `torch.where(condition, input, other)` 实现，并扩展它以支持单参数变体 `torch.where(condition)`，该变体返回非零元素的索引。

技术栈

pythonpytorch

领域

backend

议题类型

功能

难度面向新贡献者的预计实现难度，1 表示很小改动，5 表示专家级工作。

3

预计时间有经验贡献者完成调查、实现、测试并准备 pull request 的粗略时间范围。

半天

活动状态议题当前的可参与程度：新鲜、活跃、陈旧、阻塞或等待维护者输入。

活跃

清晰度议题是否清楚说明期望改动、验收标准和下一步。

清晰

前置要求

PythonPyTorch

新手友好度1-100 的估计分数，表示该议题对首次贡献者的友好程度。

60