Hang: sampling VonMises distribution gets stuck in rejection sampling for small kappa

[julian-urban]

🐛 Describe the bug

Sampling the VonMises distribution gets stuck in rejection sampling for small values of the concentration parameter. With the location parameter set to 1, the problem starts around ~1e-4 for single and ~1e-8 for double precision.

torch.distributions.von_mises.VonMises(torch.Tensor([1]), torch.Tensor([1e-4])).sample()

Numpy doesn't have this problem, likely because small values of the concentration are caught and handled explicitly, see here. Implementing this might solve the issue.

Versions

Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (conda-forge gcc 12.1.0-17) 12.1.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.8.13 (default, Oct 21 2022, 23:50:54) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-52-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: NVIDIA T600 Laptop GPU
Nvidia driver version: 515.76
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] torch==1.12.1
[pip3] torchaudio==0.12.1
[pip3] torchvision==0.13.1
[conda] blas 1.0 mkl
[conda] cudatoolkit 10.2.89 hfd86e86_1
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] mkl 2021.4.0 h06a4308_640
[conda] mkl-service 2.4.0 py38h7f8727e_0
[conda] mkl_fft 1.3.1 py38hd3c417c_0
[conda] mkl_random 1.2.2 py38h51133e4_0
[conda] numpy 1.21.5 py38h6c91a56_3
[conda] numpy-base 1.21.5 py38ha15fc14_3
[conda] pytorch 1.12.1 py3.8_cuda10.2_cudnn7.6.5_0 pytorch
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] torchaudio 0.12.1 py38_cu102 pytorch
[conda] torchvision 0.13.1 py38_cu102 pytorch

cc @ezyang @gchanan @zou3519 @fritzo @neerajprad @alicanb @nikitaved @VitalyFedyunin @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @mruberry @rgommers

[mruberry]

This example causes Colab to hang. High priority for severity

[malfet]

Reproduces in nightlies:

>>> torch.distributions.von_mises.VonMises(torch.Tensor([1]), torch.Tensor([1e-4])).sample()
^CTraceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/nshulga/git/pytorch/pytorch/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/Users/nshulga/git/pytorch/pytorch/torch/distributions/von_mises.py", line 118, in sample
    return _rejection_sample(self.loc, self.concentration, self._proposal_r, x)
  File "/Users/nshulga/git/pytorch/pytorch/torch/jit/_trace.py", line 1212, in wrapper
    return fn(*args, **kwargs)
  File "/Users/nshulga/git/pytorch/pytorch/torch/distributions/von_mises.py", line 62, in _rejection_sample
    accept = ((c * (2 - c) - u2) > 0) | ((c / u2).log() + 1 - c >= 0)