Lightning-AI / Lightning-AI/lightning-thunder
`python_print` doesn't work as expected
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
### To Reproduce
Issue 1:
```python
import thunder
import torch
from thunder.core import prims
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linears = torch.nn.ModuleList(torch.nn.Linear(1, 1) for _ in range(3))
def forward(self, x):
for l in self.linears:
prims.python_print("l")
prims.python_print(x.sum())
x = l(x)
return x
fn = MyModel()
x = torch.randn(10, 1)
fn = thunder.jit(fn, disable_torch_autograd=True)
fn(x)
```
Output:
```python
tensor(2.7485)
tensor(6.1350)
tensor(-3.8157)
```
`l` is missing
Issue 2:
```python
import thunder
import torch
from thunder.core import prims
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linears = torch.nn.ModuleList(torch.nn.Linear(1, 1) for _ in range(3))
def forward(self, x):
prims.python_print("l0")
prims.python_print(x.sum())
x = self.linears[0](x)
prims.python_print("l1")
prims.python_print(x.sum())
x = self.linears[1](x)
prims.python_print("l2")
prims.python_print(x.sum())
x = self.linears[2](x)
return x
fn = MyModel()
x = torch.randn(10, 1)
fn = thunder.jit(fn, disable_torch_autograd=True)
fn(x)
```
Output:
```python
l0
l1
l2
tensor(-1.1433)
tensor(8.6309)
tensor(-3.1484)
```
### Expected behavior
The prints should appear where they should, as many times as they should.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start at thunder.core.prims.python_print and trace how calls are captured and emitted by thunder.jit. Reproduce both examples to compare string and tensor output ordering and repetition. Done means the printed values appear in source order and execute once per loop iteration as expected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100