Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

enable using python_callable without mapping symbols to their impls #1168

Open
t-vi opened this issue Sep 18, 2024 · 1 comment
Open

enable using python_callable without mapping symbols to their impls #1168

t-vi opened this issue Sep 18, 2024 · 1 comment

Comments

@t-vi
Copy link
Collaborator

t-vi commented Sep 18, 2024

Currently, we can't use python_callable if we want to retrace because it does half of transform for execution.
We should absolutely get a flag on python_callable for that because we have a workaround in Thunder quite often now (in thunder/core/transforms.py but also in the autograd.Function handling in jit_ext now).
(the ["output"] might not make this example trivial, but still).

# Using trc.python_callable() makes it impossible to retrace the
# function because the python_callable uses python_ctx which replaces
# symbol occurrences with its symbol._call_ctx function
computation_trc = dce(computation_trc)
@wraps(computation_trc.python_callable())
def python_callable(*args, **kwargs):
return eval_trace(computation_trc, *args, **kwargs)["output"]

cc @apaz-cli

@t-vi
Copy link
Collaborator Author

t-vi commented Sep 18, 2024

This is another, more prototypical, example the commented line does not work, but it should with a flag (or maybe even by default, so only transform for execution does it differently).

@wraps(trace_of_augmented_fwd.python_callable())
def core_of_augmented_forward(*args, **kwargs):
return thunder.core.trace_interpreter.interpret_trace(trace_of_augmented_fwd, *args, **kwargs)
# core_of_augmented_forward = trace_of_augmented_fwd.python_callable(include_decorators=False)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

No branches or pull requests

1 participant