Lightning-AI / Lightning-AI/lightning-thunder

variable naming issue with parameter specification

Open
#2,555 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.