Lightning-AI/lightning-thunder

Handle returning NamedTuples from the JIT

Open

#147 opened on Apr 9, 2024

View on GitHub
 (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

NamedTuples have (some) support in the interpreter, but to return them from the JIT, we would need some more things:

  • ensure value tracking covers creation (populating the attribute_wrappers) - this is in interpreter,
  • when seeing NamedTuples being returned, add epilogue code to create the NamedTuple from its contents (which are available),

This is nontrivial as it requires return values to be routed through the epilogue, but I don't think that is grave.

This is probably an intermediate to advanced issue for people fond of the JIT / frontend bits.

Contributor guide

Research direction
Study the JIT epilogue generation and understand how NamedTuples are created in the interpreter. Then implement epilogue code to construct a NamedTuple from its tracked values.
Tech stack
pythonpytorch
Domain
backend
Issue type
Feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-2 days
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Fresh
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Clear
Prerequisites
PythonPyTorch
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
30