Make resize_ use normal device dispatch #42240
Closed
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.
resize_
only requires manual registration toAutograd
key and its device kernels can safely live together with our normal device dispatch innative_functions.yaml
.But currently we do manual registration for
CPU/CUDA
kernels (and leaves no dispatch in native_functions.yaml) which makesresize_
non-overrideable from backend point of view. While it indeed should dispatch at device level, this caused xla to whitelistresize_
and register a lowering to XLA key. This PR moves the device dispatch ofresize_
back tonative_functions.yaml
so that it shows up asabstract
method properly for downstream extensions.Note that we also do manual registration for
copy_/detach_/resize_as_/etc
in aten but they are slightly different thanresize_
since for them we only registercatchAll
kernels instead of device kernels. I'll need to investigate and send a followup PR for those ops.