Add mutated_args field to custom_op (#123129)

If provided, we:
- autogenerate an ADInplaceOrView implementation
- assume that no mutated inputs are returned as outputs. There are
  already aliasing runtime checks that check this.

Test Plan:
- new tests
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123129
Approved by: https://github.com/albanD
ghstack dependencies: #123108, #123109, #123110
This commit is contained in:
rzou
2024-04-05 11:58:23 -07:00
committed by PyTorch MergeBot
parent 9e8d2b6de2
commit 81e7a7c955
7 changed files with 151 additions and 24 deletions

View File

@ -730,6 +730,8 @@ void initDispatchBindings(PyObject* module) {
py_context_manager_DEPRECATED<at::AutoDispatchBelowAutograd>(
m, "_AutoDispatchBelowAutograd");
py_context_manager<at::AutoDispatchBelowADInplaceOrView>(
m, "_AutoDispatchBelowADInplaceOrView");
// Prints out the name of every operator that has a kernel registered to the
// Dispatcher under [dispatch_key]. If no arguments are specified, it'll print