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

[carmocca]

🚀 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 ThunderModules.

[kshitij12345]

I think another thing to consider would be if thunder.jit fuses an operation outside the module forward with operations in the forward.

def foo(module, x):
    return module(x) + 1  # it is likely that a FusionExecutor will fuse `+1` with the last functions in `module.forward`.