Port `diag` cpu from TH to ATen (#35100)
Summary:
Fixes: https://github.com/pytorch/pytorch/issues/24689
Pull Request resolved: https://github.com/pytorch/pytorch/pull/35100
Differential Revision: D20624868
Pulled By: VitalyFedyunin
fbshipit-source-id: bc436a62369aa9b6257e82051eabf5768652cf58
