Merge pull request #1653 from nouiz/gpushape

theano/sandbox/gpuarray/basic_ops.py

@@ -654,6 +654,15 @@ class GpuAlloc(HideC, Alloc):
 gpu_alloc = GpuAlloc()
 
 
+class GpuShape(HideC, tensor.Shape):
+    """
+    Implement Shape on the gpu.
+    """
+    def make_node(self, x):
+        return Apply(self, [x], [tensor.lvector()])
+gpu_shape = GpuShape()
+
+
 class GpuReshape(HideC, tensor.Reshape):
     """
     Implement Reshape on the gpu.

theano/sandbox/gpuarray/opt.py

@@ -14,7 +14,9 @@ from theano.sandbox.gpuarray.type import GpuArrayType
 from theano.sandbox.gpuarray.basic_ops import (host_from_gpu,
                                                gpu_from_host,
                                                gpu_alloc,
+                                               gpu_shape,
                                                GpuAlloc,
+                                               GpuShape,
                                                GpuReshape,
                                                GpuEye)
 from theano.sandbox.gpuarray.blas import gpu_dot22, GpuGemv, GpuGemm
@@ -282,3 +284,17 @@ def local_gpua_crossentropysoftmaxargmax1hotwithbias(node):
 @op_lifter([tensor.nnet.CrossentropySoftmax1HotWithBiasDx])
 def local_gpua_crossentropysoftmax1hotwithbiasdx(node):
     return GpuCrossentropySoftmax1HotWithBiasDx()
+
+
+@register_opt()
+@local_optimizer([tensor.Shape])
+def local_gpua_shape(node):
+    """
+    Can't use op_lifter as the output is on the GPU.
+    """
+    if isinstance(node.op, tensor.Shape):
+        x, = node.inputs
+        if x.owner and x.owner.op == host_from_gpu:
+            gpu_x, = x.owner.inputs
+            return [gpu_shape(gpu_x)]
+    return False

theano/sandbox/gpuarray/tests/test_basic_ops.py

@@ -36,7 +36,7 @@ from theano.sandbox.gpuarray.basic_ops import (host_from_gpu, gpu_from_host,
                                                gpu_alloc, gpu_from_cuda,
                                                cuda_from_gpu, HostFromGpu,
                                                GpuFromHost, GpuReshape,
-                                               GpuEye)
+                                               GpuEye, GpuShape)
 
 from theano.tests import unittest_tools as utt
 utt.seed_rng()
@@ -290,6 +290,26 @@ GpuAllocTester = makeTester(
 )
 
 
+def test_shape():
+    x = GpuArrayType(dtype='float32', broadcastable=[False, False, False])()
+    v = gpuarray.zeros((3, 4, 5), dtype='float32')
+    f = theano.function([x], x.shape)
+    topo = f.maker.fgraph.toposort()
+    assert numpy.all(f(v) == (3, 4, 5))
+    if theano.config.mode != 'FAST_COMPILE':
+        assert len(topo) == 4
+        assert isinstance(topo[0].op, T.opt.Shape_i)
+        assert isinstance(topo[1].op, T.opt.Shape_i)
+        assert isinstance(topo[2].op, T.opt.Shape_i)
+        assert isinstance(topo[3].op, T.opt.MakeVector)
+    mode = mode_with_gpu.excluding("local_shape_to_shape_i")
+    f = theano.function([x], x.shape, mode=mode)
+    topo = f.maker.fgraph.toposort()
+    assert numpy.all(f(v) == (3, 4, 5))
+    assert len(topo) == 1
+    assert isinstance(topo[0].op, GpuShape)
+
+
 class G_reshape(T_reshape):
     def shortDescription(self):
         return None