Repository: pytorch
Issue: Vectorized Jacobian and Hessian errors with ffts
Labels: module: autograd,triaged,module: vmap

I was testing out the vectorized Hessians and Jacobians for my use case, which involves FFTs, and got some errors. (Errors did not appear when I set `vectorize=False`.) Here are minimal examples of what I attempted to do and the errors that I got. 

Hessian Code:
```
import torch

torch._C._debug_only_display_vmap_fallback_warnings(True)

def func1(a):
  return torch.sum(a*torch.fft.irfft(torch.fft.rfft(a), len(a)))

a = torch.rand(3)
b = torch.autograd.functional.hessian(func1, a, create_graph=False, strict=False, vectorize=True)
```
Error messages:
`UserWarning: Batching rule not implemented for aten::_fft_c2c falling back to slow (for loop and stack) implementation (Triggered internally at  ../aten/src/ATen/BatchedFallback.cpp:75.)`

`UserWarning: Batching rule not implemented for aten::copy_ falling back to slow (for loop) implementation (Triggered internally at  ../aten/src/ATen/BatchedFallback.cpp:75.)`

`RuntimeError: vmap: aten::copy_(self, *extra_args) is not possible because there exists a Tensor 'other' in extra_args that has more elements than 'self'. This happened due to 'other' being vmapped over but 'self' not being vmapped over at level 1. Please try to use out-of-place operators instead of aten::copy_. If said operator is being called inside the PyTorch framework, please file a bug report instead.`

Similar results occur for the Jacobian.

Jacobian code:
```
import torch

torch._C._debug_only_display_vmap_fallback_warnings(True)

def func2(a):
  return a*torch.fft.irfft(torch.fft.rfft(a), len(a))

a = torch.rand(3)
b = torch.autograd.functional.jacobian(func2, a, create_graph=False, strict=False, vectorize=True)
```

Error messages:
`UserWarning: Batching rule not implemented for aten::_fft_r2c falling back to slow (for loop and stack) implementation (Triggered internally at  ../aten/src/ATen/BatchedFallback.cpp:75.)`

`UserWarning: Batching rule not implemented for aten::mul_.Scalar falling back to slow (for loop) implementation (Triggered internally at  ../aten/src/ATen/BatchedFallback.cpp:75.)`

`UserWarning: Batching rule not implemented for aten::copy_ falling back to slow (for loop) implementation (Triggered internally at  ../aten/src/ATen/BatchedFallback.cpp:75.)`

`RuntimeError: vmap: aten::copy_(self, *extra_args) is not possible because there exists a Tensor 'other' in extra_args that has more elements than 'self'. This happened due to 'other' being vmapped over but 'self' not being vmapped over at level 1. Please try to use out-of-place operators instead of aten::copy_. If said operator is being called inside the PyTorch framework, please file a bug report instead.`

### Versions

Collecting environment information...
PyTorch version: 1.10.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.5 (default, Jun  4 2021, 12:28:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.16.3-microsoft-standard-WSL2-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A


cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7