add a test

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"""