Repository: pytorch
Issue: `DataParallel` (`broadcast_coalesced`) with complex tensors yield real views
Labels: high priority,oncall: distributed,triaged,module: complex,module: data parallel

## 🐛 Bug

Using `DataParallel` on complex tensors (either parameters or inputs/outputs) yield real views. The expected behavior would be to obtain complex tensors on each replicate. Casting the views back to complex leads to an exception during `BroadcastBackward`. See the original issue (#47330) and PR (#48686), I believe the problem is that `broadcast_coalesced` calls `torch.view_as_real` but the corresponding `torch.view_as_complex` is never called.

Example (here only the model parameters are complex, but there are additional problems if the inputs/outputs are as well):
```python
import torch
import torch.nn as nn
import torch.nn.parallel

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        param = torch.rand((3, 2))
        param = torch.view_as_complex(param)
        self.param = nn.Parameter(param)  # complex, shape (3,)

    def forward(self, x):
        print(self.param.dtype, self.param.shape)   # prints "torch.float32 torch.Size([3, 2])" instead of "torch.complex64 torch.Size([3])"
        return torch.abs(x - torch.view_as_complex(self.param))  # view_as_complex necessary to have a complex tensor

model = nn.parallel.DataParallel(Model()).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

x = torch.rand(100, 3).cuda()
y = model(x)

loss = torch.mean(y ** 2)
loss.backward()
optimizer.step()  # Function BroadcastBackward returned an invalid gradient at index 0 - got [3, 2] but expected shape compatible with [3]
```

## Environment

```
PyTorch version: 1.8.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: version 2.8.12.2

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB

Nvidia driver version: 460.32.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] numpydoc==1.1.0
[pip3] torch==1.8.1+cu111
[pip3] torchaudio==0.8.1
[pip3] torchvision==0.9.1+cu111
[conda] blas                      1.0                         mkl  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.2.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.19.2           py38h54aff64_0  
[conda] numpy-base                1.19.2           py38hfa32c7d_0  
[conda] numpydoc                  1.1.0                      py_0  
[conda] torch                     1.8.1+cu111              pypi_0    pypi
[conda] torchaudio                0.8.1                    pypi_0    pypi
[conda] torchvision               0.9.1+cu111              pypi_0    pypi
```

## Workaround

A temporary workaround is to ensure all parameters and model inputs/outputs are real with an additional last dimension of size 2, and calling `torch.view_as_complex` on them at each forward. In the above example, one simply removes `param = torch.view_as_complex(param)` in `Model.__init__`.

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501 @anjali411 @dylanbespalko @mruberry @Lezcano @nikitaved @bdhirsh @jbschlosser @agolynski @mrzzd @cbalioglu