Lightning-AI / Lightning-AI/lightning-thunder
Make autocast a Transform
- 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