Lightning-AI / Lightning-AI/lightning-thunder

`python_print` doesn't work as expected

Open
#312 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.