Repository: cupy/cupy

Issue:
Optimize random/distributions.py
https://github.com/cupy/cupy/blob/faf0b0bc540ddc3e9f0c035c4d683fe09a8f88f3/cupy/random/distributions.py#L159


Discussion:

Resolving patch:
diff --git a/cupy/random/_kernels.py b/cupy/random/_kernels.py
--- a/cupy/random/_kernels.py
+++ b/cupy/random/_kernels.py
@@ -277,19 +277,6 @@
 '''
 
 
-gumbel_kernel = core.ElementwiseKernel(
-    'T x, T loc, T scale', 'T y',
-    'y = loc - log(-log(1 - x)) * scale',
-    'gumbel_kernel'
-)
-
-laplace_kernel = core.ElementwiseKernel(
-    'T x, T loc, T scale', 'T y',
-    'y = (x < 0.5)? loc + scale * log(x + x):'
-    ' loc - scale * log(2.0 - x - x)',
-    'laplace_kernel'
-)
-
 definitions = \
     [rk_state_difinition, rk_seed_definition, rk_random_definition,
      rk_double_definition, rk_binomial_btpe_definition,
diff --git a/cupy/random/distributions.py b/cupy/random/distributions.py
--- a/cupy/random/distributions.py
+++ b/cupy/random/distributions.py
@@ -182,7 +182,4 @@ def uniform(low=0.0, high=1.0, size=None, dtype=float):
 
     """
     rs = generator.get_random_state()
-    x = rs.uniform(0.0, 1.0, size=size, dtype=dtype)
-    cupy.multiply(x, (high - low), out=x)
-    cupy.add(x, low, out=x)
-    return x
+    return rs.uniform(low, high, size=size, dtype=dtype)
diff --git a/cupy/random/generator.py b/cupy/random/generator.py
--- a/cupy/random/generator.py
+++ b/cupy/random/generator.py
@@ -90,6 +90,11 @@ def binomial(self, n, p, size=None, dtype=int):
             self.rk_seed += numpy.prod(size)
         return y
 
+    _laplace_kernel = core.ElementwiseKernel(
+        'T x, T loc, T scale', 'T y',
+        'y = T(loc) + T(scale) * ((x < 0.5) ? log(x + x): -log(2.0 - x - x))',
+        'laplace_kernel')
+
     def laplace(self, loc=0.0, scale=1.0, size=None, dtype=float):
         """Returns an array of samples drawn from the laplace distribution.
 
@@ -97,10 +102,12 @@ def laplace(self, loc=0.0, scale=1.0, size=None, dtype=float):
             :func:`cupy.random.laplace` for full documentation,
             :meth:`numpy.random.RandomState.laplace`
         """
-        x = self.uniform(size=size, dtype=dtype)
-        loc = cupy.asarray(loc, dtype)
-        scale = cupy.asarray(scale, dtype)
-        _kernels.laplace_kernel(x, loc, scale, x)
+        x = self.random_sample(size=size, dtype=dtype)
+        if not numpy.isscalar(loc):
+            loc = cupy.asarray(loc, dtype)
+        if not numpy.isscalar(scale):
+            scale = cupy.asarray(scale, dtype)
+        RandomState._laplace_kernel(x, loc, scale, x)
         return x
 
     def lognormal(self, mean=0.0, sigma=1.0, size=None, dtype=float):
@@ -164,6 +171,16 @@ def randn(self, *size, **kwarg):
     _1m_kernel = core.ElementwiseKernel(
         '', 'T x', 'x = 1 - x', 'cupy_random_1_minus_x')
 
+    def _random_sample_raw(self, size, dtype):
+        dtype = _check_and_get_dtype(dtype)
+        out = cupy.empty(size, dtype=dtype)
+        if dtype.char == 'f':
+            func = curand.generateUniform
+        else:
+            func = curand.generateUniformDouble
+        func(self._generator, out.data.ptr, out.size)
+        return out
+
     def random_sample(self, size=None, dtype=float):
         """Returns an array of random values over the interval ``[0, 1)``.
 
@@ -172,13 +189,7 @@ def random_sample(self, size=None, dtype=float):
             :meth:`numpy.random.RandomState.random_sample`
 
         """
-        dtype = _check_and_get_dtype(dtype)
-        out = cupy.empty(size, dtype=dtype)
-        if dtype.char == 'f':
-            func = curand.generateUniform
-        else:
-            func = curand.generateUniformDouble
-        func(self._generator, out.data.ptr, out.size)
+        out = self._random_sample_raw(size, dtype)
         RandomState._1m_kernel(out)
         return out
 
@@ -310,6 +321,11 @@ def tomaxint(self, size=None):
         sample &= cupy.iinfo(cupy.int_).max
         return sample
 
+    _scale_kernel = core.ElementwiseKernel(
+        'T low, T high', 'T x',
+        'x = T(low) + x * T(high - low)',
+        'cupy_scale')
+
     def uniform(self, low=0.0, high=1.0, size=None, dtype=float):
         """Returns an array of uniformly-distributed samples over an interval.
 
@@ -320,7 +336,11 @@ def uniform(self, low=0.0, high=1.0, size=None, dtype=float):
         """
         dtype = numpy.dtype(dtype)
         rand = self.random_sample(size=size, dtype=dtype)
-        return dtype.type(low) + rand * dtype.type(high - low)
+        if not numpy.isscalar(low):
+            low = cupy.asarray(low, dtype)
+        if not numpy.isscalar(high):
+            high = cupy.asarray(high, dtype)
+        return RandomState._scale_kernel(low, high, rand)
 
     def choice(self, a, size=None, replace=True, p=None):
         """Returns an array of random values from a given 1-D array.
@@ -447,6 +467,11 @@ def permutation(self, num):
             array = cupy.argsort(sample)
         return array
 
+    _gumbel_kernel = core.ElementwiseKernel(
+        'T x, T loc, T scale', 'T y',
+        'y = T(loc) - log(-log(x)) * T(scale)',
+        'gumbel_kernel')
+
     def gumbel(self, loc=0.0, scale=1.0, size=None, dtype=float):
         """Returns an array of samples drawn from a Gumbel distribution.
 
@@ -454,12 +479,12 @@ def gumbel(self, loc=0.0, scale=1.0, size=None, dtype=float):
             :func:`cupy.random.gumbel` for full documentation,
             :meth:`numpy.random.RandomState.gumbel`
         """
-        x = self.uniform(size=size, dtype=dtype)
-        loc = cupy.asarray(loc, dtype)
-        scale = cupy.asarray(scale, dtype)
-        # We use `1 - x` as input of `log` method to prevent overflow.
-        # It obeys numpy implementation.
-        _kernels.gumbel_kernel(x, loc, scale, x)
+        x = self._random_sample_raw(size=size, dtype=dtype)
+        if not numpy.isscalar(loc):
+            loc = cupy.asarray(loc, dtype)
+        if not numpy.isscalar(scale):
+            scale = cupy.asarray(scale, dtype)
+        RandomState._gumbel_kernel(x, loc, scale, x)
         return x
 
     def randint(self, low, high=None, size=None, dtype='l'):


Test patch: