Richard Zou 2932b0bf37 Extend impl_backward to be usable with torch.library operators (#106817) ...

- impl_save_for_backward/impl_backward only work for functional,
non-view schemas. We validate this.
- impl_save_for_backward/impl_backward raise if there already exists an
autograd implementation from torch.library / TORCH_LIBRARY.
- Operators constructed via custom_op receive an "autograd indirection
kernel". The "autograd indirection kernel" automatically pulls the
constructed autograd kernel out of a dict. When
impl_save_for_backward/impl_backward get used with torch.library
operators, we also register the "autograd indirection kernel" so we can
reuse the logic.

Test Plan:
- new tests
Pull Request resolved: https://github.com/pytorch/pytorch/pull/106817
Approved by: https://github.com/soulitzer
ghstack dependencies: #106799, #106800

2023-08-14 14:33:46 +00:00

..

__init__.py

[custom_op] Create a new torch._custom_op namespace (#101823)

2023-05-23 18:31:29 +00:00

autograd.py

Extend impl_backward to be usable with torch.library operators (#106817)

2023-08-14 14:33:46 +00:00

functional.py

Add API to construct the functional variant of an op (#102293)

2023-06-02 13:36:50 +00:00

impl.py

Extend impl_backward to be usable with torch.library operators (#106817)

2023-08-14 14:33:46 +00:00