Merge pull request #2544 from nouiz/crash_fix

theano/misc/check_blas.py

@@ -113,7 +113,7 @@ def jobman_job(state, channel):
 
 
 def test():
-    execute()
+    return execute()
 
 
 parser = OptionParser(
@@ -222,6 +222,7 @@ if __name__ == "__main__":
         GTX 980           0.06s        
         GTX 970           0.08s
         GTX 680                         0.11s  0.12s  0.154s               0.218s
+        GRID K520         0.14s
         GTX 580                         0.16s  0.16s  0.164s               0.203s
         GTX 480                         0.19s  0.19s  0.192s               0.237s 0.27s
         GTX 750 Ti        0.20s

theano/sandbox/cuda/opt.py

@@ -310,7 +310,8 @@ def local_gpu_split(node):
             any([c != 'output' and isinstance(c.op, GpuFromHost) for c, idx
                  in outs_clients])):
             new_op = GpuSplit(node.op.len_splits)
-            split_res = new_op(gpu_from_host(input), *node.inputs[1:])
+            split_res = new_op(gpu_from_host(input), *node.inputs[1:],
+                               return_list=True)
             return [host_from_gpu(o) for o in split_res]
     return False
 

theano/sandbox/cuda/tests/test_opt.py

@@ -321,6 +321,21 @@ def test_local_gpu_split():
     # Check equality
     assert all([(cpu == gpu).all() for cpu, gpu in zip(cpu_res, gpu_res)])
 
+    # Test that split with only 1 output work
+    ra = tensor.split(x, splits, n_splits=1, axis=0)
+    f = theano.function([x, splits], [ra], mode=mode_without_gpu)
+    cpu_res = f([0, 1, 2, 3, 4, 5], [6])
+    l = f.maker.fgraph.toposort()
+    # Ensure that one op is theano.tensor.Split
+    assert any([isinstance(o.op, theano.tensor.Split) for o in l])
+    # GPU version
+    f = theano.function([x, splits], [ra], mode=mode_with_gpu)
+    gpu_res = f([0, 1, 2, 3, 4, 5], [6])
+    l = f.maker.fgraph.toposort()
+    assert any([isinstance(o.op, theano.sandbox.cuda.GpuSplit) for o in l])
+    # Check equality
+    assert all([(cpu == gpu).all() for cpu, gpu in zip(cpu_res, gpu_res)])
+
 
 def test_print_op():
     """ Test that print ops don't block gpu optimization"""