Lightning-AI / Lightning-AI/lightning-thunder
Unhelpful variable renaming
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
The commit 7c916c13675bb05b1a5522a9c797b33e997e4f19 tries to give identifiers observed in the source to the proxies of variables. However such identifiers sometimes refer to thunder's internal implementation of the language's construct, and this can be confusing.
### Code sample
```py
@thunder.jit
def f(xs, s):
for i, x in enumerate(xs):
s += x
return s
n = 6
xs = [torch.zeros(n) for _ in range(n)]
s = torch.zeros(n)
f(xs, s)
print(thunder.last_traces(f)[-1])
```
```py
@torch.no_grad()
@no_autocast
def computation(res, elem, x, b, t_0_4, t_0_5, s):
# res: "cpu f32[6]"
# elem: "cpu f32[6]"
# xs: "cpu f32[6]"
# b: "cpu f32[6]"
# t_0_4: "cpu f32[6]"
# t_0_5: "cpu f32[6]"
# s: "cpu f32[6]"
t0 = torch.add(s, res) # t0: "cpu f32[6]"
# ...
```
These renamings happen in [`thunder/core/jit_ext._maybe_update_proxy_name`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/jit_ext.py#L1108), which is called from [`thunder.core.interpreter._load_fast_handler`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L4572). `frame.code` reveals that the irrelevant variable names come from the lookasides implemented in `thunder.core.interpreter`. `res` is from [`SequenceIter.__next__`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L1997), `elem` is from [`_enumerate_lookaside`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L1382), `b` from [`_binary_op`](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/interpreter.py#L3228).
### Ideal behavior
When deciding the identifiers we can just ignore those in thunder's interpreter, which will result in
```py
def computation(x, t_0_1, t_0_2, t_0_3, t_0_4, t_0_5, s):
```
Altenatively, we can perhaps label the identifier `x` by an index when `x` is bound to multiple proxies, as
```py
def computation(x_0, x_1, x_2, x_3, x_4, t_0_5, s):
```
cc @t-vi @nikitaved
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 at thunder/core/jit_ext.py:_maybe_update_proxy_name and trace its call from thunder/core/interpreter.py:_load_fast_handler. Inspect the interpreter lookasides, including SequenceIter.__next__, _enumerate_lookaside, and _binary_op, to distinguish implementation identifiers from source identifiers. Done means generated proxy names no longer expose irrelevant interpreter names, with the intended naming behavior resolved for the provided example.
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
- Mostly clear
- Newbie friendliness
- 35/100