Lightning-AI / Lightning-AI/lightning-thunder
Updating a nn.Module attribute in forward raises an exception in prologue trace.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
```python
import torch
import thunder
import thunder.examine
class MyModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.bar = 1
def forward(self, x):
self.bar = self.bar + 1
# self.bar = 2 # This works
return x
m = MyModule()
x = torch.randn(16, 16, device='cuda')
jit_linear = thunder.jit(m)
o = jit_linear(x)
```
Error:
```
File "/home/kkalambarkar/lightning-thunder/thunder/__init__.py", line 537, in get_computation_and_inputs
inps = pro(*args, **kwargs)
File "/home/kkalambarkar/git/pytorch/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/home/kkalambarkar/miniconda3/envs/pytorch-dev/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "thunder.prologue_0", line 16, in prologue
File "/home/kkalambarkar/lightning-thunder/thunder/executors/pythonex.py", line 100, in _check_number_type_and_value_impl
utils.check(
File "/home/kkalambarkar/lightning-thunder/thunder/core/baseutils.py", line 103, in check
raise exception_type(s())
RuntimeError: Expected 2 to be equal to and have the type of 1
```
cc @apaz-cli
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 by reproducing the example in the issue, then inspect thunder/__init__.py around get_computation_and_inputs and thunder/executors/pythonex.py around _check_number_type_and_value_impl. Trace how the prologue handles self.bar changing from 1 to 2. Done means the example no longer raises the type/value error while preserving the attribute update.
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
- 35/100