Make diag_embed a primitive w.r.t. autograd (#71750)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71750
functorch has problems vmap-ing over diag_embed due to the in-place copy_.
This PR adds a backward formula for it so that it becomes a primitive
w.r.t. autograd.
Test Plan: - tested with functorch
Reviewed By: anjali411
Differential Revision: D33772753
Pulled By: zou3519
fbshipit-source-id: da8ff3a10a1de1d60e6de6292003079d4b5ba861
(cherry picked from commit afe9059bfb1f2856e463e6ae988ec0ae86fdd470)
