Lightning-AI / Lightning-AI/lightning-thunder

How do I access the `ThunderModule` if I'm compiling a function?

Open
#121 1 comment 0 reactions 0 assignees View on GitHub
enhancement help wanted
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature

### Motivation

Sometimes the code requires that a `ThunderModule` is passed, however, if the user is compiling a function that takes the module as an argument, the user doesn't have a way to get a reference to it.

For example, #96 implements a workaround for this issue with the `no_sync` context manager.

### Pitch

Provide an API to get this reference. Maybe it's something like `thunder.compile_data(jitted_function).module`.

### Additional context

The design might need to consider the presence of multiple `ThunderModule`s.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reviewing the existing no_sync workaround from issue #96 and the proposed thunder.compile_data(jitted_function) entry point. Determine how a compiled function should expose its ThunderModule, including behavior when multiple ThunderModules are present. Done means the API and its multiple-module semantics are specified and implemented.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.