Lightning-AI / Lightning-AI/lightning-thunder

Better name for elements of list in `prologue_trace` and `computation_trace`

Open
#152 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement help wanted jit
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.