Repository: pytorch
Issue: Poor-quality random numbers generated by torch.poisson on gpus
Labels: module: distributions,triaged

### 🐛 Describe the bug

Random numbers sampled from `torch.poisson` on a gpu are poor-quality, this leads easly to wrong statistical predictions. The effect is enhanced when sampling occurs in blocks of size `(512, 512)`.
The following code demonstrates this by sampling two poisson distributions and looking at their difference (for a single entry of the `(512, 512)` sized tensor). The result should be skellam distributed and hence the mean converge to the difference of the two entries and the variance to their sum.

```python
import torch as pt
import pandas as pd
import numpy as np
from typing import Tuple

def update_default(n_current: int,
                   tensor_1: pt.Tensor,
                   tensor_2: pt.Tensor,
                   M1_aggregate: pt.Tensor,
                   M2_aggregate: pt.Tensor,
                   generator : pt.Generator) -> None:
    noisy_tensor_1 = pt.poisson(tensor_1, generator=generator)
    noisy_tensor_2 = pt.poisson(tensor_2, generator=generator)

    diff_tensor = noisy_tensor_2 - noisy_tensor_1

    delta_M1 = diff_tensor - M1_aggregate
    M1_aggregate += delta_M1 / n_current
    delta_M2 = diff_tensor - M1_aggregate
    M2_aggregate += delta_M1 * delta_M2

def run(n_total: int,
        tensor_1: pt.Tensor,
        tensor_2: pt.Tensor,
        generator : pt.Generator) -> Tuple[pt.Tensor, pt.Tensor]:
    
    M1_agg = pt.full_like(tensor_1, 0.)
    M2_agg = pt.full_like(tensor_1, 0.)

    for i in range(0, n_total):
        update_default(n_current=i+1, tensor_1=tensor_1, tensor_2=tensor_2, M1_aggregate=M1_agg, M2_aggregate=M2_agg, generator=generator)

    return M1_agg, M2_agg / n_total

def get_estimates_at_index(n: int,
                           ind: list[int, int],
                           tensor1: pt.Tensor,
                           tensor2: pt.Tensor,
                           generator : pt.Generator) -> Tuple[float, float]:
    
    mean, var = run(n, tensor1, tensor2, generator)
    return mean[*ind].item(), var[*ind].item()

def get_numpy_estimate(n:int, val1:float, val2:float) -> Tuple[float, float]:
    diff = np.random.poisson(val2, (n,)).astype(float) - np.random.poisson(val1, (n,)).astype(float)
    mean = diff.mean()
    var = diff.var()

    return mean, var

device = 'cuda'
g_cuda = pt.Generator(device=device)
g_cuda.manual_seed(7324786)
np.random.default_rng(9954786)

shape = (512, 512)
center = [256, 256]

ti_1 = pt.full(shape, 21.2, device=device, dtype=float)
ti_2 = ti_1.clone()
ti_2[*center] = 1.5 * ti_2[*center]

ind_results = []
for n in [100, 1000, 10000, 100000, 1000000, 10000000]:
    mean_default, var_default = get_estimates_at_index(n, center, ti_1, ti_2, g_cuda)
    mean_np, var_np = get_numpy_estimate(n, ti_1[*center].item(), ti_2[*center].item())
    ind_results.append(pd.DataFrame({
                'n': n,
                'mean_pt': mean_default,
                'var_pt': var_default,
                'mean_np': mean_np,
                'var_np': var_np},
                index=[0]))

result = pd.concat(ind_results, ignore_index=True).sort_values(by='n')

print(result)

# https://en.wikipedia.org/wiki/Skellam_distribution
print(f"Expectation mean: {ti_2[*center] - ti_1[*center]}")
print(f"Expectation Variance : {ti_2[*center] + ti_1[*center]}")
```

The code runs for a bit (you can reduce the number of samples to obtain a faster indication), and eventually returns:
```
          n    mean_pt     var_pt    mean_np     var_np
0       100   9.980000  45.059600  10.330000  48.101100
1      1000  10.541000  50.198319  10.705000  51.733975
2     10000  10.661000  48.355879  10.683200  52.867838
3    100000  10.567200  48.061524  10.589090  52.812283
4   1000000  10.604558  48.009594  10.600764  53.057829
5  10000000  10.601304  48.071595  10.599137  52.978110
Expectation mean: 10.599999999999998
Expectation Variance : 53.0
```

Notes:
- `numpy` included for reference
- It works correctly when setting device to `cpu`
- Code verified using normal distributions and device `cpu`

### Versions

Collecting environment information...
PyTorch version: 2.4.1
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Enterprise
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.5 | packaged by Anaconda, Inc. | (main, Sep 12 2024, 18:18:29) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-11-10.0.22621-SP0
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2050
Nvidia driver version: 552.41
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=1400
DeviceID=CPU0
Family=1
L2CacheSize=18432
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=1400
Name=Intel(R) Core(TM) Ultra 7 165H
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.4.1
[pip3] torchaudio==2.4.1
[pip3] torchvision==0.19.1
[conda] blas                      1.0                         mkl
[conda] mkl                       2023.1.0         h6b88ed4_46358
[conda] mkl-service               2.4.0           py312h2bbff1b_1
[conda] mkl_fft                   1.3.10          py312h827c3e9_0
[conda] mkl_random                1.2.7           py312h0158946_0
[conda] numpy                     1.26.4          py312hfd52020_0
[conda] numpy-base                1.26.4          py312h4dde369_0
[conda] pytorch                   2.4.1           py3.12_cuda12.4_cudnn9_0    pytorch
[conda] pytorch-cuda              12.4                 h3fd98bf_6    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.4.1                    pypi_0    pypi
[conda] torchvision               0.19.1                   pypi_0    pypi

cc @fritzo @neerajprad @alicanb @nikitaved