Enable codegen of per-dispatch key derivative formulas in derivatives.yaml
derivatives.yaml can now take a dispatch entry which registers per-autograd dispatch key derivatives such as
name: foo(Tensor self, Tensor y) -> Tensor
dispatch:
default:
x: grad
y: grad.expand(y.sizes())
AutogradNestedTensor:
x: grad
y: NestedTensor_foo_backward(grad, y)
output_differentiabilty: [True]
However the old schema where there is no dispatch entry is still supported.
Would greatly appreciate feedback on how to improve the testing strategy of this PR, currently have registered an aten test op in TestOps.cpp with dummy gradients in derivatives.yaml and have some tests in test_autograd.py:TestAutogradMultipleDispatch but I am not sure whether these are sufficiently rigorous.
Additionally, this PR also makes the assumption that sets like VIEW_FUNCTIONS are per-native-function and not per-native-function-and-dispatch-key. I'm not sure whether this is necessarily the case, would there ever be a situation where (e.g. a nested_tensor op is a view op but the aten function is not or vice versa?)
Stack from ghstack:
- #82916
- -> #82801
- -> #82801
:link: Helpful links
- :test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/82801
- :page_facing_up: Preview Python docs built from this PR
- :page_facing_up: Preview C++ docs built from this PR
- :question:Need help or want to give feedback on the CI? Visit our office hours
:white_check_mark: No Failures (0 Pending)
As of commit 388e8fbd2f (more details on the Dr. CI page):
Expand to see more
:green_heart: :green_heart: Looks good so far! There are no failures yet. :green_heart: :green_heart:
This comment was automatically generated by Dr. CI (expand for details).
Please report bugs/suggestions to the (internal) Dr. CI Users group.
@pytorchbot merge
@pytorchbot successfully started a merge and created land time checks. See merge status here and land check progress here.