Fix underflow issue with dirichlet sample (pytorch#17488)

Summary:
Addresses pytorch#15738, using fritzo's suggestion. This adds a `torch._sample_dirichlet` method in `Distributions.cpp` and `Distributions.cu`.
 - For CPU, this leads to no perf hit since all we do is to promote the `alpha` to double when getting the gamma samples (the gamma sampler anyways uses `accscalar_t`(double for CPU)) and cast it back to float32 on return.
 - I have added an analogous method for CUDA as well, but the default sampler for CUDA uses scalar_t for efficiency, so I have kept it as that. With this, I do not see the bias towards 1 as reported in pytorch#15738 with `float32`, but there is a spurious mode at 0.5, as would be expected. Users would need to explicitly use `float64` for GPU to not see the spurious mode at 0.5. (EDIT: see note below, it appears that the bias issue is still there for certain builds).

Added some tests and checked that there is no perf regression. My experience with C++ is very limited, so apologies in advance if I missed something basic. cc. ailzhang, fritzo, fmassa
Pull Request resolved: pytorch#17488

Differential Revision: D14410301

Pulled By: ezyang

fbshipit-source-id: 62b2f694b4642685eab06db96d74ce28e05c3992

Author: neerajprad

aten/src/ATen/native/Distributions.cpp

@@ -228,4 +228,41 @@ Tensor _s_gamma_cpu(const Tensor& alpha, Generator *gen) {
   return ret;
 }
 
+Tensor _s_dirichlet_cpu(const Tensor& alpha, Generator *gen) {
+  Tensor ret = at::zeros(alpha.sizes(), alpha.options());
+  AT_DISPATCH_FLOATING_TYPES(ret.type(), "dirichlet", [&] {
+    Tensor gamma = at::zeros(alpha.sizes(), alpha.options().dtype(ScalarType::Double));
+    THGenerator* generator = get_generator(gen);
+    std::lock_guard<std::mutex> lock(generator->mutex);
+    /* Generate gamma sample by casting alpha to double to prevent underflow. */
+    CPU_tensor_apply2<double, scalar_t>(gamma, alpha,
+      [generator](double& ret_val, const scalar_t& alpha){
+        auto uniform_lambda = [generator] () {
+          return THRandom_standard_uniform(generator);
+        };
+        BaseSampler<double, decltype(uniform_lambda)> standard_uniform(uniform_lambda);
+
+        auto normal_lambda = [generator] () {
+          return THRandom_normal(generator, 0.0, 1.0);
+        };
+        BaseSampler<double, decltype(normal_lambda)> standard_normal(normal_lambda);
+        auto sample = sample_gamma<double, double, decltype(uniform_lambda), decltype(normal_lambda)>
+          (alpha, standard_uniform, standard_normal);
+        ret_val = std::max(std::numeric_limits<double>::min(), sample);
+      }
+    );
+    /* Normalize and cast back to scalar_t. */
+    Tensor gamma_sum = gamma.sum(-1, true).expand(alpha.sizes());
+    CPU_tensor_apply3<scalar_t, double , double>(ret, gamma, gamma_sum,
+      [](scalar_t& ret_val, const double& gamma, const double& gamma_sum){
+        ret_val = gamma / gamma_sum;
+        auto min_val = std::numeric_limits<scalar_t>::min();
+        auto max_val = std::nexttoward(static_cast<scalar_t>(1.0f), 0.0f);
+        ret_val = std::min(max_val, std::max(min_val, ret_val));
+        ret_val = static_cast<scalar_t>(ret_val);
+      }
+    );
+  });
+  return ret;
+}
 }} // namespace at::native

aten/src/ATen/native/cuda/Distributions.cu

@@ -82,7 +82,7 @@ void gamma_cuda_kernel(
         };
         BaseSampler<accscalar_t, decltype(normal_lambda)> standard_normal(normal_lambda);
         auto sample = sample_gamma<scalar_t, accscalar_t, decltype(uniform_lambda), decltype(normal_lambda)>(alpha, standard_uniform, standard_normal);
-        auto min_value = std::numeric_limits<scalar_t>::lowest();
+        auto min_value = std::numeric_limits<scalar_t>::min();
         ret_val = (min_value > sample) ? min_value : sample;
       });
 }
@@ -181,6 +181,21 @@ void bernoulli_scalar_cuda_kernel(
     );
 }
 
+template<typename scalar_t>
+void dirichlet_scalar_cuda_kernel(
+    at::Tensor& ret,
+    const at::Tensor& gamma) {
+  auto gamma_sum = gamma.sum(-1, true).expand(ret.sizes());
+  at::cuda::CUDA_tensor_apply3<scalar_t, scalar_t, scalar_t>(ret, gamma, gamma_sum,
+  [] __device__(scalar_t &ret_val, const scalar_t &gamma, const scalar_t &gamma_sum) {
+    ret_val = gamma / gamma_sum;
+    auto min_value = std::numeric_limits<scalar_t>::min();
+    auto max_value = 1 - std::numeric_limits<scalar_t>::epsilon();
+    ret_val = (min_value > ret_val) ? min_value : ret_val;
+    ret_val = (max_value < ret_val) ? max_value : ret_val;
+  });
+}
+
 } // namespace
 
 namespace at { namespace native {
@@ -200,6 +215,16 @@ Tensor _s_gamma_cuda(const Tensor& alpha, Generator* gen) {
   return ret;
 }
 
+Tensor _s_dirichlet_cuda(const Tensor& alpha, Generator* gen) {
+  Tensor ret = at::empty(alpha.sizes(), alpha.options());
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(ret.type(), "dirichlet", [&] {
+    Tensor gamma = at::empty(alpha.sizes(), alpha.options());
+    gamma_cuda_kernel<scalar_t>(gamma, alpha, next_philox_seed(gen, 10));
+    dirichlet_scalar_cuda_kernel<scalar_t>(ret, gamma);
+  });
+  return ret;
+}
+
 Tensor _standard_gamma_grad_cuda(const Tensor& self, const Tensor& output) {
   Tensor ret = at::empty(self.sizes(), self.options());
   AT_DISPATCH_FLOATING_TYPES_AND_HALF(self.scalar_type(), "_standard_gamma_grad_cuda", [&] {

aten/src/ATen/native/native_functions.yaml

@@ -2455,6 +2455,13 @@
     CPU: _s_gamma_cpu
     CUDA: _s_gamma_cuda
 
+- func: _sample_dirichlet(Tensor self, Generator? generator=None) -> Tensor
+  matches_jit_signature: True
+  variants: function
+  dispatch:
+    CPU: _s_dirichlet_cpu
+    CUDA: _s_dirichlet_cuda
+
 - func: poisson(Tensor self, Generator? generator=None) -> Tensor
   matches_jit_signature: True
   dispatch:

test/test_distributions.py

@@ -2208,6 +2208,39 @@ def test_beta_sample(self):
             x = Beta(Tensor([1e-6]), Tensor([1e-6])).sample()[0]
             self.assertTrue(np.isfinite(x) and x > 0, 'Invalid Beta.sample(): {}'.format(x))
 
+    def test_beta_underflow(self):
+        # For low values of (alpha, beta), the gamma samples can underflow
+        # with float32 and result in a spurious mode at 0.5. To prevent this,
+        # torch._sample_dirichlet works with double precision for intermediate
+        # calculations.
+        set_rng_seed(1)
+        num_samples = 50000
+        for dtype in [torch.float, torch.double]:
+            conc = torch.tensor(1e-2, dtype=dtype)
+            beta_samples = Beta(conc, conc).sample([num_samples])
+            self.assertEqual((beta_samples == 0).sum(), 0)
+            self.assertEqual((beta_samples == 1).sum(), 0)
+            # assert support is concentrated around 0 and 1
+            frac_zeros = float((beta_samples < 0.1).sum()) / num_samples
+            frac_ones = float((beta_samples > 0.9).sum()) / num_samples
+            self.assertEqual(frac_zeros, 0.5, 0.05)
+            self.assertEqual(frac_ones, 0.5, 0.05)
+
+    @unittest.skipIf(not TEST_CUDA, "CUDA not found")
+    def test_beta_underflow_gpu(self):
+        set_rng_seed(1)
+        num_samples = 50000
+        conc = torch.tensor(1e-2, dtype=torch.float64).cuda()
+        beta_samples = Beta(conc, conc).sample([num_samples])
+        self.assertEqual((beta_samples == 0).sum(), 0)
+        self.assertEqual((beta_samples == 1).sum(), 0)
+        # assert support is concentrated around 0 and 1
+        frac_zeros = float((beta_samples < 0.1).sum()) / num_samples
+        frac_ones = float((beta_samples > 0.9).sum()) / num_samples
+        # TODO: increase precision once imbalance on GPU is fixed.
+        self.assertEqual(frac_zeros, 0.5, 0.12)
+        self.assertEqual(frac_ones, 0.5, 0.12)
+
     def test_independent_shape(self):
         for Dist, params in EXAMPLES:
             for param in params:
@@ -3375,7 +3408,7 @@ def test_entropy_monte_carlo(self):
                     continue
                 x = dist.sample(sample_shape=(60000,))
                 expected = -dist.log_prob(x).mean(0)
-                ignore = (expected == inf)
+                ignore = (expected == inf) | (expected == -inf)
                 expected[ignore] = actual[ignore]
                 self.assertEqual(actual, expected, prec=0.2, message='\n'.join([
                     '{} example {}/{}, incorrect .entropy().'.format(Dist.__name__, i + 1, len(params)),

torch/distributions/dirichlet.py

@@ -1,17 +1,8 @@
-from numbers import Number
-
 import torch
 from torch.autograd import Function
 from torch.autograd.function import once_differentiable
 from torch.distributions import constraints
 from torch.distributions.exp_family import ExponentialFamily
-from torch.distributions.utils import broadcast_all, clamp_probs
-
-
-def _dirichlet_sample_nograd(concentration):
-    probs = torch._standard_gamma(concentration)
-    probs /= probs.sum(-1, True)
-    return clamp_probs(probs)
 
 
 # This helper is exposed for testing.
@@ -24,7 +15,7 @@ def _Dirichlet_backward(x, concentration, grad_output):
 class _Dirichlet(Function):
     @staticmethod
     def forward(ctx, concentration):
-        x = _dirichlet_sample_nograd(concentration)
+        x = torch._sample_dirichlet(concentration)
         ctx.save_for_backward(x, concentration)
         return x
 
@@ -71,9 +62,7 @@ def expand(self, batch_shape, _instance=None):
     def rsample(self, sample_shape=()):
         shape = self._extended_shape(sample_shape)
         concentration = self.concentration.expand(shape)
-        if isinstance(concentration, torch.Tensor):
-            return _Dirichlet.apply(concentration)
-        return _dirichlet_sample_nograd(concentration)
+        return _Dirichlet.apply(concentration)
 
     def log_prob(self, value):
         if self._validate_args: