Repository: pytorch
Issue: [Feature Request] An alternative sampling routine for Dirichlet to fix Dirichlet and Beta sampling bugs
Labels: module: distributions,feature,triaged

### 🚀 The feature, motivation and pitch

I'm a long-time user but a first-time contributor, and this is my first issue. I propose the implementation of the log-gamma distribution into Pytorch and to use it to increase precision in the sampling of Dirichlet/Beta distributions.

This feature proposal is in part related to [issue #84625](https://github.com/pytorch/pytorch/issues/84625). This issue reveals that for small parameter values, the beta distribution behaves incorrectly. More generally, the Dirichlet distribution behaves incorrectly for small parameter values. 

With small parameter values (below 1e-3), the Dirichlet distribution gives a constant vector. This is not the expected behaviour. We would expect to see a vector with one entry equal to 1 and the others equal to 0 (that is, a coordinate on the N-simplex).

The reason why this occurs is because of the way we generate samples of the Dirichlet distribution. Pytorch, like Numpy, samples a Dirichlet distribution by using gamma variables (see [Distributions.cpp](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Distributions.cpp)) This involves dividing a gamma variable by a sum of gamma variables. By doing a manual example of this we can see how the error arises

```
import numpy as np
concentration = [1e-5, 1e-5, 1e-5, 1e-5]
gamma_vars = [torch.distributions.Gamma(p, 1) for p in concentration]
print("Gamma variable output is incredibly small: ", gamma_vars[1].sample())
samples = [gamma_var.sample() for gamma_var in gamma_vars]
vector = [sample / np.sum(samples) for sample in samples]
print("Now calculating the Dirichlet distributed vector gives us", vector)
print("Since the Gamma variables are so small, we end up getting a constant output here")
print("This is what causes us to observe the behaviour of very large values of the parameter when we input very small values")
```

This error doesn't occur in Numpy, because Numpy truncates very small gamma values to 0, meaning that instead of calculating something like `1e-38/(4 * 1e-38) = 1/4`, it just outputs `NaN` due to division by 0.

A possible fix would be to:
- Use logarithms to avoid excessively small output values from gamma. Sample in log-space to achieve better precision for small parameter values. This would involve writing a routine to sample from the log-gamma distribution. This is exactly [what Tensorflow does ](https://github.com/tensorflow/probability/blob/dc5b6b094cfc1f18f949e120c3491acb2e2af6eb/tensorflow_probability/python/distributions/gamma.py#L453-L459)
- Add a "hack" for small values to force them to be 0/output NaN like Numpy does. This means that gamma won't output such exceptionally small values anymore.

**My feature request** is to implement my first suggestion, to implement the log-gamma distribution into Pytorch, and use it to fix this bug.

### Alternatives

Numpy's solution is to allow gamma to generate infinitely small values, in which case it outputs NaN. See the following example.

Alternatively (or additionally), could add a warning to the documentation.

### Additional context

I would be happy to work on solving this when the issue has been discussed further.

cc @fritzo @neerajprad @alicanb @nikitaved