Richard Zou 3897c479af Add API to construct the functional variant of an op (#102293) ...

`register_functional_op`:
- constructs the functional variant of an op
- registers a functionalization kernel to the op

To get this to work:
- `register_functional_op` makes assumptions that it checks about the
op's schema. In particular, the op is not allowed to return anything it
mutates. We can relax these constraints in the future.
- We add a "boxed" python functionalization kernel that handles this
case.

I'm not actually sure (or convinced) this should be public API or how
it should work. If we want this to be public, then it should probably be
a torch.library API, but does that also mean we should give the same
lifetime guarantees? If so, then it would be up to the user to construct
a Library object to actually register the functional variant onto.

Test Plan:
- new tests
Pull Request resolved: https://github.com/pytorch/pytorch/pull/102293
Approved by: https://github.com/bdhirsh

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

..

__init__.py

…

autograd.py

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

2023-06-02 13:36:50 +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

[custom_op] explicit autograd API (#101824)

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