Lightning-AI/lightning-thunder
Better name for elements of list in `prologue_trace` and `computation_trace`
Aberta
#152 aberto em 10 de abr. de 2024
enhancementhelp wantedjit
Métricas do repositório
- Stars
- (1.460 estrelas)
- Métricas de merge de PR
- (Nenhuma PRs mesclada em 30d)
Description
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.
Prologue Trace
# 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
# 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]