This repository was archived by the owner on Aug 21, 2025. It is now read-only.
Add trace decomposition, invoke from C++ #740
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
eellison
changed the title
eellison
commented
Apr 27, 2022
Chillee
approved these changes
Apr 27, 2022
functorch/csrc/BatchRulesViews.cpp
Outdated
| return std::make_tuple(at::_unsafe_view(self_, view_size), 0); | ||
| } | ||
|
|
||
| Tensor trace_decomp(const Tensor& self) { |
There was a problem hiding this comment.
can we make this a macro? or perhaps a boxed fallback?
There was a problem hiding this comment.
Done... depends on pytorch/pytorch#76493
eellison
commented
Apr 27, 2022
added 2 commits
April 27, 2022 15:00
zou3519
reviewed
Apr 28, 2022
functorch/_src/vmap.py
Outdated
Comment on lines
40
to
47
| def register_jit_decomposition(decomp): | ||
| assert decomp in decomposition_table, f"could not find {decomp}" | ||
| decomp_fn = decomposition_table[decomp] | ||
| scripted_decomp_fn = torch.jit.script(decomp_fn) | ||
| torch.jit._register_decomposition(decomp, scripted_decomp_fn.graph) | ||
|
|
||
|
|
||
| register_jit_decomposition(torch.ops.aten.trace.default) |
There was a problem hiding this comment.
Decompositions aren't a vmap-specific thing, we also would like to apply them to our other transforms (jvp for now). Maybe a better place to put these lines would be eager_transforms.py ?
zou3519
approved these changes
Apr 28, 2022
|
@zou3519 couldn't you have also just called into the Python directly |
|
@ezyang, we have three short-term options really:
There's an open question around what the medium and long term state of this should be. That I am not sure about.
|
added 2 commits
May 2, 2022 10:33
|
I think the appropriate long term attitude is that decomps get traced through in a JIT (e.g., AOTAutograd via torchdynamo) and therefore it doesn't matter that they directly get run in Python. |
zou3519
pushed a commit
to zou3519/pytorch
that referenced
this pull request
Jul 20, 2022
…ch#740) * [WIP] add trace example * use better api * Move decomp registration * fix flake * oop * use boxed fallback * rename * move and fix string handling Co-authored-by: Elias Ellison <eellison@devfair044.h1.fair>
bigfootjon
pushed a commit
to pytorch/pytorch
that referenced
this pull request
Jul 21, 2022
…ch#740) * [WIP] add trace example * use better api * Move decomp registration * fix flake * oop * use boxed fallback * rename * move and fix string handling Co-authored-by: Elias Ellison <eellison@devfair044.h1.fair>
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
TODO - integrate registrations of the decomposition w/jit into the
decompositions.pyfileRelies on pytorch/pytorch#76252
cc @zou3519 @Chillee