Make trace introspection available as methods or attribute accesses on functions returned by Thunder

[IvanYashchuk]

🚀 Feature

Make trace introspection available as methods or attribute accesses on functions returned by thunder.jit.

Looking at the last traces of a thunder callable currently requires importing Thunder explicitly. They are currently hidden behind ._lc_cs attribute: https://github.com/Lightning-AI/lightning-thunder/blob/a0a61518a70d80ab5ce4d1b6e44bcae159d81bb0/thunder/init.py#L836

[t-vi]

We should not add alternative ways to do things much, it is good when there is only one way to achieve things. What's bad with using thunder.last_traces?

[IvanYashchuk]

There's nothing bad with thunder.last_traces, the alternative is better:

@thunder.jit
def f():
    pass

# Call once ...
f()
# ... and access the traces with
f.last_traces

It doesn't matter in scripts but it's more friendly to REPL development and exploration. Users can hit TAB and see what can be done with this object that thunder.jit returns. With thunder.last_traces as a free function there's an additional effort involved of knowing and understanding which objects are allowed to be passed to thunder.last_traces.

[tfogal]

+1. if we don't want multiple ways to do things, maybe we just remove the thunder.last_traces function, such that f.last_traces is the only way to do this

[mruberry]

+1. if we don't want multiple ways to do things, maybe we just remove the thunder.last_traces function, such that f.last_traces is the only way to do this

I think that's a popular function -- we'd probably have to deprecate it if we want to change the UX.

[tfogal]

maybe we just remove the thunder.last_traces function, such that f.last_traces is the only way to do this

that's a popular function -- we'd probably have to deprecate it

fair point!

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.