Merge pull request #2092 from nouiz/gpu_assert

theano/sandbox/cuda/opt.py

@@ -1680,6 +1680,16 @@ def local_gpualloc(node):
         return [new_out]
 
 
+@register_opt()
+@local_optimizer([theano.tensor.opt.Assert])
+def local_assert(node):
+    if (isinstance(node.op, theano.tensor.opt.Assert) and
+        node.inputs[0].owner and
+        isinstance(node.inputs[0].owner.op,
+                   HostFromGpu)):
+        return [host_from_gpu(node.op(node.inputs[0].owner.inputs[0]))]
+
+
 @register_opt()
 @local_optimizer([GpuAlloc])
 def local_gpualloc_memset_0(node):

theano/sandbox/cuda/tests/test_opt.py

@@ -41,6 +41,17 @@ def test_no_shared_var_graph():
     assert numpy.any(isinstance(x.op,cuda.GpuFromHost) for x in l)
     assert numpy.any(isinstance(x.op,cuda.HostFromGpu) for x in l)
 
+
+def test_local_assert():
+    x = theano.tensor.fmatrix()
+    a = theano.tensor.opt.assert_op(x, theano.tensor.eq(x, 0).any())
+    f = theano.function([x], a, mode=mode_with_gpu)
+    topo = f.maker.fgraph.toposort()
+    a_op = [n for n in topo if isinstance(n.op, theano.tensor.opt.Assert)]
+    assert len(a_op) == 1
+    assert isinstance(a_op[0].inputs[0].type, CudaNdarrayType)
+
+
 def test_int_pow():
     a = CudaNdarrayType([False])()
 

theano/sandbox/gpuarray/opt.py

@@ -496,6 +496,16 @@ def local_gpua_softmaxwithbias(node):
     return GpuSoftmaxWithBias()
 
 
+@register_opt()
+@local_optimizer([theano.tensor.opt.Assert])
+def local_assert(node):
+    if (isinstance(node.op, theano.tensor.opt.Assert) and
+        node.inputs[0].owner and
+        isinstance(node.inputs[0].owner.op,
+                   HostFromGpu)):
+        return [host_from_gpu(node.op(node.inputs[0].owner.inputs[0]))]
+
+
 @register_opt()
 @op_lifter([gpu_from_host, ConvOp])
 def local_gpu_conv(node):

theano/sandbox/gpuarray/tests/test_opt.py

@@ -16,6 +16,16 @@ from theano.tests.unittest_tools import SkipTest
 from theano.tensor.tests.test_basic import TestSpecifyShape
 
 
+def test_local_assert():
+    x = theano.tensor.fmatrix()
+    a = theano.tensor.opt.assert_op(x, theano.tensor.eq(x, 0).any())
+    f = theano.function([x], a, mode=mode_with_gpu)
+    topo = f.maker.fgraph.toposort()
+    a_op = [n for n in topo if isinstance(n.op, theano.tensor.opt.Assert)]
+    assert len(a_op) == 1
+    assert isinstance(a_op[0].inputs[0].type, GpuArrayType)
+
+
 def test_flatten():
     m = theano.tensor.fmatrix()
     f = theano.function([m], m.flatten(), mode=mode_with_gpu)