Lightning-AI / Lightning-AI/lightning-thunder
Better name for elements of list in `prologue_trace` and `computation_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 thunder
import torch
def foo(xs):
result = []
for x in xs:
result.append(x + x)
return result
jfoo = thunder.jit(foo)
o = jfoo([torch.randn(3,),] * 6)
print(thunder.last_prologue_traces(jfoo)[-1])
print(thunder.last_traces(jfoo)[-1])
```
Names for the arguments to the computation trace are : `res, x, a, b, t_0_4, t_0_5`. It would be nice if there was a consistent pattern.
Traces
Prologue Trace
```python
# Constructed by Transform for execution (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast()
def prologue(*args, **kwargs):
# args: "Any"
check_len(args, 1)
# prims.check_len(args, 1)
# kwargs: "Any"
check_len(kwargs, 0)
# prims.check_len(kwargs, 0)
subscr: "Any" = args[0]
res: "cpu f32[3]" = subscr[0]
x: "cpu f32[3]" = subscr[1]
a: "cpu f32[3]" = subscr[2]
b: "cpu f32[3]" = subscr[3]
t_0_4: "cpu f32[3]" = subscr[4]
t_0_5: "cpu f32[3]" = subscr[5]
...
return (res, x, a, b, t_0_4, t_0_5)
```
Computation Trace
```python
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast()
def computation(res, x, a, b, t_0_4, t_0_5):
# res: "cpu f32[3]"
# x: "cpu f32[3]"
# a: "cpu f32[3]"
# b: "cpu f32[3]"
# t_0_4: "cpu f32[3]"
# t_0_5: "cpu f32[3]"
result = torch.add(res, res) # result: "cpu f32[3]"
# result = ltorch.add(res, res, alpha=None) # result: "cpu f32[3]"
# result = prims.add(res, res) # result: "cpu f32[3]"
del res
...
return [result, t1, t2, t3, t4, t5]
```
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 with the example using thunder.jit, last_prologue_traces, and last_traces, then trace where the computation-trace argument names are generated. Define a consistent naming pattern for list elements and confirm that the prologue and computation traces use it consistently.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100