Lightning-AI / Lightning-AI/lightning-thunder

`where`: inconsistency between eager PyTorch and NVFuser

Open
#833 0 comments 0 reactions 0 assignees View on GitHub
bug operators
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

We have a dtype mismatch when it comes to dealing with non-tensor objects:

```python
In [1]: import torch, thunder

In [2]: def f(x, y):
...: return torch.where(x > 0, -y, y)
...:

In [3]: jf = thunder.jit(f)

In [4]: x = torch.rand(3, 3, device="cuda")

In [5]: f(x, 0.).dtype
Out[5]: torch.float32

In [6]: jf(x, 0.).dtype
Out[6]: torch.float64

In [7]: jf(x.cpu(), 0.).dtype
Out[7]: torch.float32

...

In [10]: jf(x, 0.).dtype
Out[10]: torch.float64

In [11]: trace = thunder.last_traces(jf)[-1]

In [12]: trace
Out[12]:
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x):
# x: "cuda:0 f32[3, 3]"
[t2] = nvFusion0(x)
# t1 = prims.gt(x, 0.0) # t1: "cuda:0 b8[3, 3]"
# t2 = prims.where(t1, -0.0, 0.0) # t2: "cuda:0 f32[3, 3]"
del x
return t2

```

cc @apaz-cli

Contributor guide

No contributing guide indexed for this repository

Research direction

Run the provided Python reproduction with torch.where, thunder.jit, and CUDA tensors, then inspect thunder.last_traces(jf)[-1] and the generated nvFusion0 computation. Compare eager PyTorch and JIT results for the scalar non-tensor operand; done means the CUDA JIT path produces the same dtype as eager PyTorch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.