Merge pull request #1086 from nouiz/specify_shape

theano/tensor/basic.py

@@ -2302,6 +2302,8 @@ class SpecifyShape(Op):
 
     @note:     Maybe in the future we will never do the assert!
     @note:     We currently don't support specifying partial shape information.
+
+    @todo:     test this op with sparse and cuda ndarray. Do c code for them too.
     """
     view_map = {0: [0]}
 
@@ -2318,11 +2320,16 @@ class SpecifyShape(Op):
         if not isinstance(x, Variable):
             x = as_tensor_variable(x)
         shape = as_tensor_variable(shape)
+        assert shape.ndim == 1
+        assert "int" in shape.dtype
+        if isinstance(shape, TensorConstant):
+            assert shape.data.size == x.ndim
         return Apply(self, [x, shape], [x.type()])
 
     def perform(self, node, inp, out_):
         x, shape = inp
         out, = out_
+        assert x.ndim == shape.size
         assert numpy.all(x.shape == shape), ("got shape", x.shape,
                                            "expected", shape)
         out[0] = x
@@ -2362,6 +2369,47 @@ class SpecifyShape(Op):
             return [None]
         return self.make_node(eval_points[0], *inputs[1:]).outputs
 
+    def c_code(self, node, nodename, inp, out, sub):
+        if not isinstance(node.inputs[0], TensorVariable):
+            # The c code bellow support only Tensor.  super.c_code
+            # will raise an exception to tell that there isn't c code
+            # for the other cases.
+            return super(SpecifyShape, self).c_code(node, nodename,
+                                                    inp, out, sub)
+        iname, shape = inp
+        oname, = out
+        fail = sub['fail']
+
+        return """
+        if (PyArray_NDIM(%(iname)s) != PyArray_DIMS(%(shape)s)[0]) {
+            PyErr_Format(PyExc_AssertionError,
+                         "SpecifyShape: vector of shape have %%d element,"
+                         " but the input have %%d dimensions.",
+                         PyArray_NDIM(%(iname)s),
+                         PyArray_DIMS(%(shape)s)[0]);
+            %(fail)s;
+        }
+        for(int i = 0; i < PyArray_NDIM(%(iname)s); i++){
+            dtype_%(shape)s shp = ((dtype_%(shape)s*)PyArray_GETPTR1(%(shape)s,
+                                                                     i))[0];
+            if (PyArray_DIMS(%(iname)s)[i] != shp) {
+                PyErr_Format(PyExc_AssertionError,
+                             "SpecifyShape: dim %%d of input have shape %%d,"
+                             " expected %%d.",
+                             i, PyArray_DIMS(%(iname)s)[i],
+                             shp);
+                %(fail)s;
+            }
+        }
+        Py_XDECREF(%(oname)s);
+        %(oname)s = %(iname)s;
+        Py_XINCREF(%(oname)s);
+        """ % locals()
+
+    def c_code_cache_version(self):
+        return (1,)
+
+
 specify_shape = SpecifyShape()
 
 

theano/tensor/tests/test_basic.py

@@ -6326,6 +6326,54 @@ def test_transpose():
     assert tensor.transpose(tensor.dmatrix()).name is None
 
 
+class TestSpecifyShape(unittest.TestCase):
+    def shortDescription(self):
+        return None
+
+    def test_bad_shape(self):
+        """ Test that at run time we raise an exception when the shape
+        is not the one specified"""
+        specify_shape = SpecifyShape()
+
+        x = vector()
+        xval = numpy.random.rand(2).astype(floatX)
+        f = theano.function([x], specify_shape(x, [2]))
+        f(xval)
+        xval = numpy.random.rand(3).astype(floatX)
+        self.assertRaises(AssertionError, f, xval)
+
+        x = matrix()
+        xval = numpy.random.rand(2, 3).astype(floatX)
+        f = theano.function([x], specify_shape(x, [2, 3]))
+        f(xval)
+        for shape in [(1, 3), (2, 2), (5, 5)]:
+            xval = numpy.random.rand(*shape).astype(floatX)
+            self.assertRaises(AssertionError, f, xval)
+
+    def test_bad_number_of_shape(self):
+        """ Test that the number of dimensions provided is good"""
+        specify_shape = SpecifyShape()
+
+        x = vector()
+        shape_vec = ivector()
+        xval = numpy.random.rand(2).astype(floatX)
+        self.assertRaises(AssertionError, specify_shape, x, [])
+        self.assertRaises(AssertionError, specify_shape, x, [2, 2])
+
+        f = theano.function([x, shape_vec], specify_shape(x, shape_vec))
+        self.assertRaises(AssertionError, f, xval, [])
+        self.assertRaises(AssertionError, f, xval, [2, 2])
+
+        x = matrix()
+        xval = numpy.random.rand(2, 3).astype(floatX)
+        for shape in [(),
+                      (1,),
+                      (2, 3, 4)]:
+            self.assertRaises(AssertionError, specify_shape, x, shape)
+            f = theano.function([x, shape_vec], specify_shape(x, shape_vec))
+            self.assertRaises(AssertionError, f, xval, shape)
+
+
 class TestInferShape(utt.InferShapeTester):
 
     def test_infer_shape(self):