fixed datatype issues

theano/sandbox/multinomial.py

@@ -12,6 +12,7 @@ if cuda_available:
     from theano.sandbox.cuda import CudaNdarrayType
     from theano.sandbox.cuda.basic_ops import host_from_gpu, gpu_from_host
     from theano.sandbox.cuda.opt import register_opt
+import sys
 
 
 class MultinomialFromUniform(Op):
@@ -218,7 +219,10 @@ class WeightedSelectionFromUniform(Op):
         if unis.ndim != 1:
             raise NotImplementedError('unis ndim should be 1', unis.ndim)
         if self.odtype == 'auto':
-            odtype = pvals.dtype
+            if sys.maxsize > 2**32:
+                odtype = 'int64'
+            else:
+                odtype = 'int32'
         else:
             odtype = self.odtype
         out = T.tensor(dtype=odtype, broadcastable=pvals.type.broadcastable)
@@ -241,8 +245,12 @@ class WeightedSelectionFromUniform(Op):
             raise ValueError("unis.shape[0] != pvals.shape[0] * n_samples",
                              unis.shape[0], pvals.shape[0], n_samples)
 
+        if sys.maxsize > 2**32:
+            odtype = 'int64'
+        else:
+            odtype = 'int32'
         if z[0] is None or not numpy.all(z[0].shape == [pvals.shape[0], n_samples]):
-            z[0] = -1 * numpy.ones((pvals.shape[0], n_samples), dtype='int')
+            z[0] = -1 * numpy.ones((pvals.shape[0], n_samples), dtype=odtype)
 
         nb_multi = pvals.shape[0]
         nb_outcomes = pvals.shape[1]

theano/sandbox/tests/test_weighted_select.py

@@ -19,6 +19,7 @@ class test_OP(unittest.TestCase):
         f = function([p, u, n], m, allow_input_downcast=True)
 
         n_elements = 1000
+        all_indices = range(n_elements)
         numpy.random.seed(12345)
         for i in [5, 10, 50, 100, 500, n_elements]:
             uni = numpy.random.rand(i).astype(config.floatX)
@@ -27,7 +28,7 @@ class test_OP(unittest.TestCase):
             res = f(pvals, uni, i)
             res = numpy.squeeze(res)
             assert len(res) == i
-            assert numpy.all(numpy.in1d(numpy.unique(res), res)), res
+            assert numpy.all(numpy.in1d(numpy.unique(res), all_indices)), res
 
     def test_fail_select_alot(self):
         """
@@ -63,13 +64,13 @@ class test_OP(unittest.TestCase):
 
         n_elements = 100
         n_selected = 10
-        mean_rtol = 0.04
+        mean_rtol = 0.0005
         numpy.random.seed(12345)
         pvals = numpy.random.randint(1, 100, (1, n_elements)).astype(config.floatX)
         pvals /= pvals.sum(1)
         avg_pvals = numpy.zeros((n_elements,))
 
-        for rep in range(1000):
+        for rep in range(10000):
             uni = numpy.random.rand(n_selected).astype(config.floatX)
             res = f(pvals, uni, n_selected)
             res = numpy.squeeze(res)
@@ -93,6 +94,7 @@ class test_function(unittest.TestCase):
         f = function([p, n], m, allow_input_downcast=True)
 
         n_elements = 1000
+        all_indices = range(n_elements)
         numpy.random.seed(12345)
         for i in [5, 10, 50, 100, 500, n_elements]:
             pvals = numpy.random.randint(1, 100, (1, n_elements)).astype(config.floatX)
@@ -100,7 +102,7 @@ class test_function(unittest.TestCase):
             res = f(pvals, i)
             res = numpy.squeeze(res)
             assert len(res) == i
-            assert numpy.all(numpy.in1d(numpy.unique(res), res)), res
+            assert numpy.all(numpy.in1d(numpy.unique(res), all_indices)), res
 
     def test_fail_select_alot(self):
         """
@@ -137,13 +139,13 @@ class test_function(unittest.TestCase):
 
         n_elements = 100
         n_selected = 10
-        mean_rtol = 0.04
+        mean_rtol = 0.0005
         numpy.random.seed(12345)
         pvals = numpy.random.randint(1, 100, (1, n_elements)).astype(config.floatX)
         pvals /= pvals.sum(1)
         avg_pvals = numpy.zeros((n_elements,))
 
-        for rep in range(1000):
+        for rep in range(10000):
             res = f(pvals, n_selected)
             res = numpy.squeeze(res)
             avg_pvals[res] += 1