Repository: pytorch/pytorch

Issue:
Incomplete documentation for `torch.randn` on complex dtypes
### 📚 The doc issue

https://pytorch.org/docs/stable/generated/torch.randn.html does not clearly specify behavior on complex data types `torch.complex64` and `torch.complex128`. 

The current notation $\textrm{out}_{i} \sim \mathcal{N} (0, 1)$ may lead to users thinking that both the real and imaginary parts are sampled from a unit Gaussian, but in fact they are each i.i.d. sampled from $\mathcal{N} (0, 1/2)$, leading to the complex signal being (correctly, for a large class of applications) sampled from the complex random normal distribution.

The following basic snippet validates this:

```
import torch

x = torch.randn(1000, 1000, dtype=torch.complex64)
print('Complex signal power using native complex: ', torch.var(x))
print('Real part signal power using native complex: ', torch.var(torch.real(x)))
print('Imag part signal power using native complex: ', torch.var(torch.imag(x)))
```

It will produce:
```
Complex signal power using native complex:  tensor(0.9998)
Real part signal power using native complex:  tensor(0.5005)
Imag part signal power using native complex:  tensor(0.4993)
```

The same holds for `torch.complex128`.

### Suggest a potential alternative/fix

The documentation for `torch.randn` could be updated to include something similar to the following clarifications:


- For complex `dtype`, the tensor is i.i.d. sampled from a complex normal distribution with zero mean and unit variance as: $\textrm{out}_{i} \sim \mathcal{CN} (0, 1) $


- $\mathcal{CN}$ is the standard notation for the complex normal distribution (https://en.wikipedia.org/wiki/Complex_normal_distribution).

cc @svekars @brycebortree @ezyang @anjali411 @dylanbespalko @mruberry @Lezcano @nikitaved


Discussion:
Thanks for the well researched issue report. Do you think you could send us the doc PR?

Resolving patch:
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -9217,8 +9217,22 @@ def merge_dicts(*dicts):
 .. math::
     \text{{out}}_{{i}} \sim \mathcal{{N}}(0, 1)
 
+For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and
+unit variance as
+
+.. math::
+    \text{{out}}_{{i}} \sim \mathcal{{CN}}(0, 1)
+
+This is equivalent to separately sampling the real :math:`(\operatorname{{Re}})` and imaginary
+:math:`(\operatorname{{Im}})` part of :math:`\text{{out}}_i` as
+
+.. math::
+    \operatorname{{Re}}(\text{{out}}_{{i}}) \sim \mathcal{{N}}(0, \frac{{1}}{{2}}),\quad
+    \operatorname{{Im}}(\text{{out}}_{{i}}) \sim \mathcal{{N}}(0, \frac{{1}}{{2}})
+
 The shape of the tensor is defined by the variable argument :attr:`size`.
 
+
 Args:
     size (int...): a sequence of integers defining the shape of the output tensor.
         Can be a variable number of arguments or a collection like a list or tuple.
@@ -9239,6 +9253,8 @@ def merge_dicts(*dicts):
     >>> torch.randn(2, 3)
     tensor([[ 1.5954,  2.8929, -1.0923],
             [ 1.1719, -0.4709, -0.1996]])
+
+.. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution
 """.format(
         **factory_common_args
     ),
@@ -9250,8 +9266,8 @@ def merge_dicts(*dicts):
 randn_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor
 
 Returns a tensor with the same size as :attr:`input` that is filled with
-random numbers from a normal distribution with mean 0 and variance 1.
-``torch.randn_like(input)`` is equivalent to
+random numbers from a normal distribution with mean 0 and variance 1. Please refer to :func:`torch.randn` for the
+sampling process of complex dtypes. ``torch.randn_like(input)`` is equivalent to
 ``torch.randn(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``.
 
 Args:


Test patch: