Support out argument in torch.fft ops (#49335)
Summary:
Ref https://github.com/pytorch/pytorch/issues/42175
This adds out argument support to all functions in the `torch.fft` namespace except for `fftshift` and `ifftshift` because they rely on `at::roll` which doesn't have an out argument version.
Note that there's no general way to do the transforms directly into the output since both cufft and mkl-fft only support single batch dimensions. At a minimum, the output may need to be re-strided which I don't think is expected from `out` arguments normally. So, on cpu this just copies the result into the out tensor. On cuda, the normalization is changed to call `at::mul_out` instead of an inplace multiply.
If it's desirable, I could add a special case to transform into the output when `out.numel() == 0` since there's no expectation to preserve the strides in that case anyway. But that would lead to the slightly odd situation where `out` having the correct shape follows a different code path from `out.resize_(0)`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49335
Reviewed By: mrshenli
Differential Revision: D25756635
Pulled By: mruberry
fbshipit-source-id: d29843f024942443c8857139a2abdde09affd7d6
