Skip to content

Remove qualified access to methods from autodiff #2147

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.

Sign up for GitHub

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

Merged
t-vi merged 15 commits into from
May 28, 2025
Merged

Remove qualified access to methods from autodiff #2147

Changes from 3 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
094a9e7
update typing to fit the rest of thunder codebase
riccardofelluga May 27, 2025
381f003
add insert aliases in higher order functions traces
riccardofelluga May 27, 2025
e5f269e
refactor and improve typing annotation
riccardofelluga May 27, 2025
8cf0c50
joint trace outputs the output of the grad transform
riccardofelluga May 27, 2025
524142b
comment out special treatment of copy
riccardofelluga May 27, 2025
274452e
uncomment copy special handling
riccardofelluga May 27, 2025
a429bc9
insert alias updates for thunderfx path too
riccardofelluga May 27, 2025
78ff43c
alias fwd trace fx path earlier
riccardofelluga May 27, 2025
9b9068a
add insert alias tests
riccardofelluga May 27, 2025
2751579
Revert "update typing to fit the rest of thunder codebase" and adds b…
riccardofelluga May 28, 2025
5174fff
revert forward_and_backward_from_trace qualified access
riccardofelluga May 28, 2025
0493ed8
update typing to conform to the rest of thunder code base
riccardofelluga May 28, 2025
21f342a
update missing
riccardofelluga May 28, 2025
76e4eb2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 28, 2025
887ab5e
Merge branch 'joint-trace-dataflow-fix' into remove-qualified-access-…
t-vi May 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions thunder/transforms/autodiff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,9 +520,8 @@ def backward_fn(saved_for_backward, cotangents):

def forward_and_backward_from_trace(trace: TraceCtx, torch_autograd=False) -> ForwardBackwardTraces:
if not torch_autograd:
from thunder.core.transforms import forward_and_backward_from_trace as legacy_autograd
return thunder.core.transforms.forward_and_backward_from_trace(trace, torch_autograd=torch_autograd)

return legacy_autograd(trace, torch_autograd=torch_autograd)
joint_trace = grad_transform_on_trace(trace)

forward_trace, backward_trace = split_into_forward_and_backward(joint_trace)
Expand Down
Loading