Repository: pytorch/pytorch

Issue:
Implement random_gamma() sampler (with gradients)
Many probability distributions require sampling from the [Gamma distribution](https://en.wikipedia.org/wiki/Gamma_distribution), including: `Gamma`, `Beta`, and `Dirichlet`.

Since Gamma samplers have complex control flow (for rejection sampling) and are seldom a bottleneck in probabilistic algorithms, it should suffice to implement a CPU-only implementation at first. What is more important than a CUDA implementation is a reparameterized sampler so that stochastic gradients can be propagated through the sampler (see [paper](http://proceedings.mlr.press/v54/naesseth17a/naesseth17a.pdf) and [reference implementation](https://github.com/blei-lab/ars-reparameterization) by @naesseth).

## Tasks
- [x] #3841 CPU implementation of basic `random_gamma(requires_grad=False)`
- [ ] #3978 Support reparameterized `random_gamma(requires_grad=True)`

## Map of modifications
- `aten/src/TH/THRandom.c/h` random single numbers
- `aten/src/TH/generic/THTensorRandom.c/h` random tensors
- `aten/src/ATen/Declarations.cwrap` bindings for ATen
- `torch/csrc/generic/methods/TensorRandom.cwrap` bindings for torch.Tensor
- `torch/autograd/variable.py` - Variable
- `torch/distributions.py` - Distributions


Discussion:

Resolving patch:
diff --git a/torch/distributions.py b/torch/distributions.py
--- a/torch/distributions.py
+++ b/torch/distributions.py
@@ -32,8 +32,8 @@
 import math
 from numbers import Number
 import torch
-from torch.autograd import Variable
-
+from torch.autograd import Function, Variable
+from torch.autograd.function import once_differentiable
 
 __all__ = ['Distribution', 'Bernoulli', 'Categorical', 'Normal', 'Gamma']
 
@@ -198,10 +198,25 @@ def log_prob(self, value):
         return -((value - self.mean) ** 2) / (2 * var) - log_std - math.log(math.sqrt(2 * math.pi))
 
 
+class _StandardGamma(Function):
+    @staticmethod
+    def forward(ctx, alpha):
+        x = torch._C._standard_gamma(alpha)
+        ctx.save_for_backward(x, alpha)
+        return x
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        x, alpha = ctx.saved_tensors
+        grad = torch._C._standard_gamma_grad(x, alpha)
+        return grad_output * grad
+
+
 def _standard_gamma(alpha):
     if not isinstance(alpha, Variable):
         return torch._C._standard_gamma(alpha)
-    return Variable(torch._C._standard_gamma(alpha.data))
+    return _StandardGamma.apply(alpha)
 
 
 class Gamma(Distribution):
@@ -219,6 +234,7 @@ class Gamma(Distribution):
         alpha (float or Tensor or Variable): shape parameter of the distribution
         beta (float or Tensor or Variable): rate = 1 / scale of the distribution
     """
+    has_rsample = True
 
     def __init__(self, alpha, beta):
         # TODO handle (Variable, Number) cases


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -43,7 +43,10 @@ def _check_log_prob(self, dist, asset_fn):
     def _check_sampler_sampler(self, torch_dist, ref_dist, message,
                                num_samples=10000, failure_rate=1e-3):
         # Checks that the .sample() method matches a reference function.
-        torch_samples = torch_dist.sample_n(num_samples).squeeze().cpu().numpy()
+        torch_samples = torch_dist.sample_n(num_samples).squeeze()
+        if isinstance(torch_samples, Variable):
+            torch_samples = torch_samples.data
+        torch_samples = torch_samples.cpu().numpy()
         ref_samples = ref_dist.rvs(num_samples)
         samples = [(x, +1) for x in torch_samples] + [(x, -1) for x in ref_samples]
         samples.sort()
@@ -160,9 +163,38 @@ def test_gamma_sample(self):
         self._set_rng_seed()
         for alpha, beta in product([0.1, 1.0, 5.0], [0.1, 1.0, 10.0]):
             self._check_sampler_sampler(Gamma(alpha, beta),
-                                        scipy.stats.gamma(alpha, scale=1 / beta),
+                                        scipy.stats.gamma(alpha, scale=1.0 / beta),
                                         'Gamma(alpha={}, beta={})'.format(alpha, beta))
 
+    # This is a randomized test.
+    @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    def test_gamma_sample_grad(self):
+        self._set_rng_seed(1)
+        num_samples = 100
+        for alpha in [1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e2, 1e3, 1e4]:
+            alphas = Variable(torch.Tensor([alpha] * num_samples), requires_grad=True)
+            betas = Variable(torch.ones(num_samples))
+            x = Gamma(alphas, betas).sample()
+            x.sum().backward()
+            x, ind = x.data.sort()
+            x = x.numpy()
+            actual_grad = alphas.grad.data[ind].numpy()
+            # Compare with expected gradient dx/dalpha along constant cdf(x,alpha).
+            cdf = scipy.stats.gamma.cdf
+            pdf = scipy.stats.gamma.pdf
+            eps = 0.02 * alpha if alpha < 100 else 0.02 * alpha ** 0.5
+            cdf_alpha = (cdf(x, alpha + eps) - cdf(x, alpha - eps)) / (2 * eps)
+            cdf_x = pdf(x, alpha)
+            expected_grad = -cdf_alpha / cdf_x
+            rel_error = np.abs(actual_grad - expected_grad) / (expected_grad + 1e-100)
+            self.assertLess(np.max(rel_error), 0.005,
+                            '\n'.join(['Bad gradients for Gamma({}, 1)'.format(alpha),
+                                       'x {}'.format(x),
+                                       'expected {}'.format(expected_grad),
+                                       'actual {}'.format(actual_grad),
+                                       'rel error {}'.format(rel_error),
+                                       'max error {}'.format(rel_error.max())]))
+
 
 if __name__ == '__main__':
     run_tests()