added numpy.random.choice function to theano.tensor.raw_random random

theano/tensor/raw_random.py

@@ -572,6 +572,49 @@ def random_integers(random_state, size=None, low=0, high=1, ndim=None,
     return op(random_state, size, low, high)
 
 
+def choice_helper(random_state, a, replace, p, size):
+    """
+    Helper function to draw random numbers using numpy's choice function.
+
+    This is a generalization of numpy.random.choice to the case where `a`,
+    `replace` and `p` are tensors.
+    """
+    if a.ndim > 1:
+        raise ValueError('a.ndim (%i) must be 0 or 1' % a.ndim)
+    if p.ndim == 1:
+        if p.size == 0:
+            p = None
+    else:
+        raise ValueError('p.ndim (%i) must be 1' % p.ndim)
+    replace = bool(replace)
+    return random_state.choice(a, size, replace, p)
+
+
+def choice(random_state, size=None, a=2, replace=True, p=None, ndim=None,
+           dtype='int64'):
+    """
+    Choose values from `a` with or without replacement. `a` can be a 1-D array
+    or a positive scalar. If `a` is a scalar, the samples are drawn from the
+    range 0,...,a-1.
+
+    If the size argument is ambiguous on the number of dimensions, ndim
+    may be a plain integer to supplement the missing information.
+
+    If size is None, a scalar will be returned.
+    """
+    a = tensor.as_tensor_variable(a)
+    if isinstance(replace, bool):
+        replace = tensor.constant(replace, dtype='int8')
+    else:
+        replace = tensor.as_tensor_variable(replace)
+    # encode p=None as an empty vector
+    p = tensor.as_tensor_variable(p or [])
+    ndim, size, bcast = _infer_ndim_bcast(ndim, size)
+    op = RandomFunction(choice_helper, tensor.TensorType(dtype=dtype,
+                                                         broadcastable=bcast))
+    return op(random_state, size, a, replace, p)
+
+
 def permutation_helper(random_state, n, shape):
     """Helper function to generate permutations from integers.
 
@@ -829,6 +872,19 @@ class RandomStreamsBase(object):
         """
         return self.gen(random_integers, size, low, high, ndim=ndim,
                         dtype=dtype)
+    
+    def choice(self, size=None, a=2, replace=True, p=None, ndim=None,
+               dtype='int64'):
+        """
+        Choose values from `a` with or without replacement. `a` can be a 1-D
+        array or a positive scalar. If `a` is a scalar, the samples are drawn
+        from the range 0,...,a-1.
+
+        If the size argument is ambiguous on the number of dimensions,
+        ndim may be a plain integer to supplement the missing
+        information.
+        """
+        return self.gen(choice, size, a, replace, p, ndim=ndim, dtype=dtype)
 
     def permutation(self, size=None, n=1, ndim=None, dtype='int64'):
         """

theano/tensor/tests/test_raw_random.py

@@ -474,6 +474,33 @@ class T_random_function(utt.InferShapeTester):
                     update=post_r2, mutable=True)],
                 [out2], accept_inplace=True)
         self.assertRaises(ValueError, f2)
+    
+    def test_choice(self):
+        """Test that raw_random.choice generates the same
+        results as numpy."""
+        # Check over two calls to see if the random state is correctly updated.
+        rng_R = random_state_type()
+        # Use non-default parameters, and larger dimensions because of
+        # the integer nature of the result
+        post_r, out = choice(rng_R, (11, 8), 10, 1, 0)
+
+        f = compile.function(
+                [compile.In(rng_R,
+                    value=numpy.random.RandomState(utt.fetch_seed()),
+                    update=post_r, mutable=True)],
+                [out], accept_inplace=True)
+
+        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
+        val0 = f()
+        val1 = f()
+        numpy_val0 = numpy_rng.choice(10, (11, 8), True, None)
+        numpy_val1 = numpy_rng.choice(10, (11, 8), True, None)
+        print val0
+        print numpy_val0
+        print val1
+        print numpy_val1
+        self.assertTrue(numpy.allclose(val0, numpy_val0))
+        self.assertTrue(numpy.allclose(val1, numpy_val1))
 
     def test_permutation(self):
         """Test that raw_random.permutation generates the same