Composite compliance for index_copy, index_fill, masked_scatter, masked_fill (#71751)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71751
Before this PR, each of the above operations were composite and had
in-place variants that were primitive w.r.t. autograd.
The backward passes are not composite compliant due to the op (e.g.
index_copy) decomposing into index_copy_ and index_copy_'s backward
formula having in-place operations in it. To fix this, for each of the
ops mentioned in the title:
- I deleted the autograd formula for the in-place variant and replaced
it with the out-of-place variant
- This makes the forward-ad formula slightly slower because the codegen
generates a temporary but offline discussion suggests it's not worth
maintaining two sets of formulas for this and we can make the autograd
codegen smarter in the future.
- I then replaced instances of grad.clone().inplace_variant_(...) with
grad.outplace_variant(...)
Test Plan:
- run existing tests to check correctness
- run functorch tests
Reviewed By: anjali411
Differential Revision: D33772756
Pulled By: zou3519
fbshipit-source-id: fd22fe1d542e6e2a16af0865c2ddce0e65c04d70
(cherry picked from commit d025ba03270d53e19b2e68e8dd7ae49f2bb84532)
