Refactor `autocast` to `AutocastTransform` for Improved Composability and Transform Integration

[rittik9]

Before submitting

• [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• [x] Did you read the contributor guideline, Pull Request section?
• [ ] Did you make sure to update the docs?
• [ ] Did you write any new necessary tests?

What does this PR do?

Fixes https://github.com/Lightning-AI/lightning-thunder/issues/1122

PR review

Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[t-vi]

Hi @rittik9 , thank you for working on this!

So two quick comments and let me know how much you want to go into details or explore yourself:

• this should work without changing thunder/__init__.py, quite likely, the best thing is to define the transform_traces_pre_prologue method on the transform and do the changing there (see eg quantization for an example),
• it would be cool if we could move from eval_trace to using the brand-new thunder.trace_interpreter.TraceSubstitutionProcessor.

[rittik9]

@t-vi Thank you for your suggestions! I’m still getting familiar with the codebase, so it might take me some time. I’ll definitely review the examples you mentioned and work on incorporating your suggestions in this PR.

[rittik9]

@t-vi apologies for the delay, I think this one is ready for review. Can I have your thoughts... I'll add the tests soon.