Lightning-AI / Lightning-AI/lightning-thunder
variable naming issue with parameter specification
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.
## 🐛 Bug
When a `Class` is `thunder.jit`ed, if the class has an instance variable with the same name as a parameter in its forward function, and that parameter is specified by name, Thunder names them the same thing. This clearly leads to undesirable behavior.
### To Reproduce
Steps to reproduce the behavior:
#### Code sample
```
import thunder
import torch
class ExampleModel(torch.nn.Module):
def __init__(
self,
var: torch.Tensor,
):
super().__init__()
self.var = var
def forward(
self,
var: torch.Tensor,
):
element_lookup = torch.bucketize(var, self.var)
return element_lookup
base_model = ExampleModel(torch.tensor([1, 2, 3]))
thunder_model = thunder.jit(base_model)
base_lookup = base_model.forward(var=torch.Tensor([1, 3]))
thunder_lookup = thunder_model.forward(var=torch.Tensor([1, 3])) # works as expected if `var=` is not specified
print(f"base_lookup = {base_lookup}, thunder_lookup = {thunder_lookup}")
print(thunder.last_traces(thunder_model)[-1])
```
### Expected behavior
Expected output
```
base_lookup = tensor([0, 2]), thunder_lookup = tensor([0, 2])
# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(var, t_var):
# var: "cpu f32[2]"
# t_var: "cpu i64[3]"
# /teamspace/studios/this_studio/achira-optimization/examples/thunder_example.py:16: element_lookup = torch.bucketize(var, self.var)
t2 = torch.bucketize(var, t_var) # t2: "cpu i64[2]"
return (t2,)
```
Actual output
```
base_lookup = tensor([0, 2]), thunder_lookup = tensor([0, 1])
# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(var):
# var: "cpu f32[2]"
# /teamspace/studios/this_studio/achira-optimization/examples/thunder_example.py:16: element_lookup = torch.bucketize(var, self.var)
t2 = torch.bucketize(var, var) # t2: "cpu i64[2]"
return (t2,)
```
In this case, `self.var` ends up completely overwritten by `var`, and Thunder changes program behavior.
### Environment
- PyTorch Version (e.g., 1.0): `2.8.0+cu128`
- OS (e.g., Linux): `Ubuntu 24.04.3 LTS`
- How you installed PyTorch (`conda`, `pip`, source): n/a
- Build command you used (if compiling from source): n/a
- Python version: `3.10.10`
- CUDA/cuDNN version: `CUDA 12.8`
- GPU models and configuration: n/a
- Any other relevant information: n/a
### Additional context
Example works as expected if `var=` is not specified.
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 the thunder.jit entry point and reproduce the ExampleModel case using a named forward parameter, then inspect thunder.last_traces(thunder_model) to follow how the instance variable and parameter are named. Done means the generated computation keeps them distinct and the jitted result matches the eager result when var= is used.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100