Lightning-AI / Lightning-AI/lightning-thunder
Thunder's vjp doesn't compute the result for variadic arguments
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
```py
import torch
import thunder
def f(*args):
a, b = args
c = a + b
return c
def ff(*args):
return thunder.core.transforms.value_and_grad(f)(*args)
a = torch.ones(3)
b = torch.ones(3)
cff = thunder.jit(ff)
print(cff(a, b)) # Returns (tensor([2., 2., 2.]), (None, None))
```
There should be ones instead of Nones.
It's probably happening because "unpacking" operations are ignored in https://github.com/Lightning-AI/lightning-thunder/blob/f7b2e1590fd0d98985131c54653b2a95ba4729ed/thunder/core/transforms.py#L2638
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with thunder/core/transforms.py around line 2638 and run the supplied PyTorch repro for a variadic function passed to value_and_grad and thunder.jit. Done means the result contains tensors of ones instead of (None, None), with coverage for this variadic-argument case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100