Attempt fixes for AT_CUDA_ENABLED changes

Author: rachtsingh

aten/src/ATen/SharedDist.cu

@@ -0,0 +1,24 @@
+#include "ATen/ATen.h"
+#include "ATen/TensorUtils.h"
+#include "ATen/NativeFunctions.h"
+#include "ATen/Dispatch.h"
+#include "ATen/Config.h"
+
+#include <nvfunctional>
+ 
+namespace at {
+  namespace native {
+    namespace dist {
+      template<typename precision_t>
+      struct baseSampler {
+        nvstd::function<precision_t(void)> sampler;
+        baseSampler(nvstd::function<precision_t(void)> sampler): sampler(sampler) {}
+        precision_t sample() {
+          return sampler();
+        }
+      };
+    }
+  }
+}
+
+// this version is only linked if CUDA is enabled, so we can safely just use CUDA features here

aten/src/ATen/native/Distributions.cpp

@@ -1,14 +1,15 @@
 #include "ATen/ATen.h"
 #include "ATen/CPUApplyUtils.h"
 #include "ATen/Dispatch.h"
+#include "ATen/Config.h"
 #include "ATen/ExpandUtils.h"
 #include "ATen/NativeFunctions.h"
 
 #include "ATen/CPUGenerator.h"
 #include "ATen/CheckGenerator.h"
 #include "ATen/Generator.h"
 
-#include <ATen/native/Distributions.cuh>
+#include <functional>
 
 #include "TH/THRandom.h"
 
@@ -121,12 +122,23 @@ Tensor _standard_gamma_grad_cuda(const Tensor& self, const Tensor& output) {
 
 /*
  * This section is a counterpart to Distributions.cu
- *
  */
 
 namespace dist {
-  // The function `sample_poisson`
-  // is adapted from Numpy's distributions.c implementation.
+
+#if !AT_CUDA_ENABLED()
+  template<typename precision_t>
+  struct baseSampler {
+    std::function<precision_t(void)> sampler;
+    baseSampler(std::function<precision_t(void)> sampler): sampler(sampler) {}
+    precision_t sample() {
+      return sampler();
+    }
+  };
+#endif
+  
+  // The functions `sample_poisson`, `sample_gamma`
+  // are adapted from Numpy's distributions.c implementation.
   // It is MIT licensed, so here is the copyright:
 
   /* Copyright 2005 Robert Kern ([email protected])
@@ -151,6 +163,41 @@ namespace dist {
    * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
    */
 
+  
+  template<typename precision_t>
+#if AT_CUDA_ENABLED()
+  __host__ __device__
+#endif
+  precision_t sample_gamma(precision_t alpha, baseSampler<precision_t>& standard_uniform, baseSampler<precision_t>& standard_normal) {
+
+    precision_t scale = 1.0;
+
+		// Boost alpha for higher acceptance probability.
+		if (alpha < 1.0) {
+			scale *= ::pow(1 - standard_uniform.sample(), 1.0 / alpha);
+			alpha += 1.0;
+		}
+
+		// This implements the acceptance-rejection method of Marsaglia and Tsang (2000)
+		// doi:10.1145/358407.358414
+		const precision_t d = alpha - 1.0 / 3.0;
+		const precision_t c = 1.0 / ::sqrt(9.0 * d);
+		for (;;) {
+			precision_t x, y;
+			do {
+				x = standard_normal.sample();
+				y = 1.0 + c * x;
+			} while (y <= 0);
+			const precision_t v = y * y * y;
+			const precision_t u = 1 - standard_uniform.sample();
+			const precision_t xx = x * x;
+			if (u < 1.0 - 0.0331 * xx * xx)
+				return scale * d * v;
+			if (::log(u) < 0.5 * xx + d * (1.0 - v + ::log(v)))
+				return scale * d * v;
+		}
+	}
+
   THGenerator * get_generator(Generator *gen) {
     auto default_gen = &at::globalContext().defaultGenerator(Backend::CPU);
     auto gen_ = check_generator<CPUGenerator>(gen, default_gen);

aten/src/ATen/native/Distributions.cuh

@@ -1,5 +1,8 @@
+#include "ATen/ATen.h"
+#include "ATen/TensorUtils.h"
 #include "ATen/NativeFunctions.h"
 #include "ATen/Dispatch.h"
+#include "ATen/Config.h"
 #include "ATen/cuda/CUDAApplyUtils.cuh"
 #include <curand.h>
 #include <curand_kernel.h>
@@ -8,7 +11,8 @@
 #include <functional>
 #include <nvfunctional>
 
-#include "ATen/native/Distributions.cuh"
+#include "ATen/SharedDist.cu"
+#include "ATen/native/Distributions.cpp"
 
 #include <TH/THAtomic.h>