Make {select,slice,diagonal}_backward primitives wrt autograd (#64933)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64933
Fixes https://github.com/facebookresearch/functorch/issues/108
This is a short-term fix. A longer-term fix would be to either:
1. have proper {select,slice,diagonal}_embed functions
2. have efficient {select,slice,diagonal}_scatter functions (and
efficient zero tensors).
NB: I didn't use diag_embed because diag_embed is slightly different
from diagonal_backward.
There are no BC concerns because TorchScript (luckily) does not
serialize the backwards graph.
Test Plan:
- run tests
- run benchmarks.
https://gist.github.com/zou3519/e7c0774d1ac97f32aa02ec44d81e60e1.
Surprisingly the instruction count goes down. This is probably because
we create fewer autograd nodes now.
Reviewed By: ezyang
Differential Revision: D30909333
Pulled By: zou3519
fbshipit-source-id: 3b33e13010ba13b4d487b346aa9bee8a0e8c378c
