Lightning-AI / Lightning-AI/lightning-thunder

Make autocast a Transform

Open
#1,122 2 comments 0 reactions 0 assignees View on GitHub
enhancement good first issue
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

Currently autocast (`transforms/autocast.py`) is carried over from the old-style transforms on callable and invoked like

```
jit(autocast(fn))
```

it would be good to make this a `AutocastTransform` to be invoked as
```
jit(fn, transforms=[Autocast()])
```
for better composability.

This should be vaguely doable with the transforms e.g. as explained in [Thunder Sessions Episode 7](https://www.youtube.com/watch?v=gW7nLRev5eU) plus the bsym mapping provided by the existing mechanism, so it might be a good first or second issue.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with transforms/autocast.py and compare its current callable-based behavior with the transform mechanism described in Thunder Sessions Episode 7. Trace the existing bsym mapping support and transform composition entry points; done means autocast can be passed as Autocast() through jit(fn, transforms=[...]) while preserving composability.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 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.