Lightning-AI/lightning-thunder

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

Open

#152 opened on Apr 10, 2024

 (0 comments) (0 reactions) (0 assignees)Python (114 forks)github user discovery
enhancementhelp wantedjit

Repository metrics

Stars
 (1,460 stars)
PR merge metrics
 (No merged PRs in 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]

Contributor guide