[AutoDiff] Canonicalize JVP/VJP types to re-enable LoadableByAddress.#27298
Closed
dan-zheng wants to merge 9 commits intotensorflowfrom
Closed
[AutoDiff] Canonicalize JVP/VJP types to re-enable LoadableByAddress.#27298dan-zheng wants to merge 9 commits intotensorflowfrom
dan-zheng wants to merge 9 commits intotensorflowfrom
Conversation
Canonicalize JVPs/VJPs to return maximally abstracted linear map functions with `@in_guaranteed` parameters and an `@out` result. This is a necessary step towards re-enabling LoadableByAddress: linear map type is no longer computed based on the original function type. See TF-11 for more info regarding LoadableByAddress. See TF-625 for more info regarding maximally abstracted linear maps.
- Change differential generation to use only tangent buffers.
- Change pullback visitors to use only adjoint buffers.
- Change adjoint of active value propagation to use only adjoint buffers.
- Mark all tangent/adjoint value helpers as `[[deprecated]]`.
- They are not deleted because helpers may become useful after
SIL opaque values are introduced.
Will re-add in follow-up for separation of concerns.
Canonicalize JVP/VJP types so the returned original result is always indirect. This makes it so JVP/VJP type no longer depends on the abstraction of the original function type, making it possible to re-enable LoadableByAddress. Resolves TF-626.
6dcf239 to
04dca63
Compare
Contributor
Author
|
A less invasive approach for re-enabling LoadableByAddress was found: #27923 Maximal indirection will not be pursued further: the differentiation transform should not be forced to generate maximally indirect code to workaround issues caused by LoadableByAddress. |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Canonicalize JVP/VJP types so the returned original result is always indirect.
This makes it so JVP/VJP type no longer depends on the abstraction of the
original function type, making it possible to re-enable LoadableByAddress.
Resolves TF-626.
Note: this PR depends on #26709.