Merge pull request #778 from larseeri/shape_opt

theano/tensor/opt.py

@@ -533,6 +533,9 @@ class MakeVector(T.Op):
             # assume that out has correct dtype. there is no cheap way to check
             out[0][...] = inputs
 
+    def infer_shape(self, node, ishapes):
+        return [(len(ishapes),)]
+            
     def grad(self, inputs, output_gradients):
         # If the output is of an integer dtype, no gradient shall pass
         if 'int' in self.dtype:
@@ -634,6 +637,9 @@ class Shape_i(T.Op):
             #      Do not continue this madness.
             return super(Shape_i, self).c_code(node, name, (x,), (out,), sub)
 
+    def infer_shape(self, node, input_shapes):
+        return [()]
+
     def grad(self, inp, grads):
         return [None]
 

theano/tensor/tests/test_opt.py

@@ -29,6 +29,8 @@ from theano.tensor.opt import (
         mul_canonizer,
         out2in,
         Shape_i,
+        Assert,
+        MakeVector
         )
 from theano import tensor
 from theano import tensor as T
@@ -2392,7 +2394,11 @@ class test_shapeoptimizer(unittest.TestCase):
         print f([[1, 2], [2, 3]])
 
 
-class test_assert(unittest.TestCase):
+class test_assert(utt.InferShapeTester):
+
+    def setUp(self):
+        super(test_assert, self).setUp()
+    
     def test0(self):
         x=T.scalar()
         y=T.scalar()
@@ -2448,6 +2454,23 @@ class test_assert(unittest.TestCase):
         assert len(topo[0].inputs)==3
         assert topo[1].op==theano.compile.function_module.deep_copy_op
 
+    def test_infer_shape(self):
+
+        adscal = dscalar()
+        bdscal = dscalar()
+        adscal_val = numpy.random.rand()
+        bdscal_val = numpy.random.rand() + 1
+        out = theano.tensor.opt.assert_(adscal, bdscal)
+        self._compile_and_check([adscal, bdscal], [out],
+                        [adscal_val, bdscal_val], Assert)
+
+        admat = dmatrix()
+        admat_val = numpy.random.rand(3, 4)
+        adscal_val += 1
+        out = theano.tensor.opt.assert_(admat, adscal, bdscal)
+        self._compile_and_check([admat, adscal, bdscal], [out],
+                        [admat_val, adscal_val, bdscal_val], Assert)
+
         
 def test_local_mul_specialize():
     mode = theano.config.mode
@@ -3414,7 +3437,12 @@ class T_local_sum_dimshuffle(unittest.TestCase):
     # test_local_sum_divprod_dimshuffle ((a * b) / (c * d))
 
 
-def test_make_vector():
+class TestMakeVector(utt.InferShapeTester):
+
+    def setUp(self):
+        super(TestMakeVector, self).setUp()
+
+    def test_make_vector():
         b = T.bscalar()
         i = T.iscalar()
         d = T.dscalar()
@@ -3496,6 +3524,32 @@ def test_make_vector():
             except AssertionError:
                 pass
 
+    def test_infer_shape(self):
+        adscal = dscalar()
+        bdscal = dscalar()
+        aiscal = iscalar()
+        biscal = iscalar()
+        ciscal = iscalar()
+        discal = iscalar()
+        adscal_val = numpy.random.rand()
+        bdscal_val = numpy.random.rand()
+        aiscal_val = numpy.random.randint(10)
+        biscal_val = numpy.random.randint(10)
+        ciscal_val = numpy.random.randint(10)
+        discal_val = numpy.random.randint(10)
+        self._compile_and_check([adscal, aiscal],
+                            [MakeVector('float64')(adscal, aiscal)],
+                            [adscal_val, aiscal_val], MakeVector)
+
+        self._compile_and_check([adscal, bdscal, aiscal],
+                            [MakeVector('float64')(adscal, bdscal, aiscal)],
+                            [adscal_val, bdscal_val, aiscal_val], MakeVector)
+
+        self._compile_and_check([aiscal, biscal, ciscal, discal],
+                    [MakeVector('int32')(aiscal, biscal, ciscal, discal)],
+                    [aiscal_val, biscal_val, ciscal_val, discal_val],
+                     MakeVector)
+
 
 def test_local_join_1():
     #test for vector
@@ -3680,6 +3734,45 @@ def test_local_upcast_elemwise_constant_inputs():
     f = function([s], [tensor.grad(x, s)])
     f([-42, -2.1, -1, -0.5, 0, 0.2, 1, 2, 12])
 
+
+class TestShape_i(utt.InferShapeTester):
+
+    def setUp(self):
+        super(TestShape_i, self).setUp()
+
+    def test_perform(self):
+
+        advec = dvector()
+        advec_val = numpy.random.rand(3)
+        f = function([advec], Shape_i(0)(advec))
+        out = f(advec_val)
+        assert numpy.allclose(out, advec_val.shape[0])
+
+        admat = dmatrix()
+        admat_val = numpy.random.rand(4, 3)
+        for i in xrange(2):
+            f = function([admat], Shape_i(i)(admat))
+            out = f(admat_val)
+            assert numpy.allclose(out, admat_val.shape[i])
+
+    def test_infer_shape(self):
+        admat = dmatrix()
+        admat_val = numpy.random.rand(3, 4)
+        self._compile_and_check([admat], [Shape_i(0)(admat)],
+                        [admat_val], Shape_i)
+
+        self._compile_and_check([admat], [Shape_i(1)(admat)],
+                        [admat_val], Shape_i)
+
+
 if __name__ == '__main__':
+    t = TestMakeVector('setUp')
+    t.setUp()
+    #t.test_perform()
+    t.test_infer_shape()
+
+    """
 #    unittest.main()
     test_fusion().tes_memory_leak()
+    """
+