Lightning-AI / Lightning-AI/lightning-thunder
`where`: inconsistency between eager PyTorch and NVFuser
- 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