Merge pull request #6357 from gvtulder/f-gpuarray-spatialtf-shape

theano/gpuarray/c_code/dnn_sptf_grid.c

@@ -53,15 +53,6 @@ APPLY_SPECIFIC(dnn_sptf_grid)(PyGpuArrayObject * theta,
         return 1;
     }
 
-    if ( PyGpuArray_DIM( theta, 1 ) != 2 || PyGpuArray_DIM( theta, 2 ) != 3 )
-    {
-        PyErr_Format( PyExc_RuntimeError,
-            "GpuDnnTransformerGrid: incorrect dimensions for theta, expected (%d, %d, %d), got (%d, %d, %d)",
-            PyGpuArray_DIMS( theta )[0], 2, 3, PyGpuArray_DIMS( theta )[0],
-            PyGpuArray_DIMS( theta )[1], PyGpuArray_DIMS( theta )[2] );
-        return 1;
-    }
-
     if ( PyArray_NDIM( out_dims ) != 1 || PyArray_SIZE( out_dims ) != 4 )
     {
         PyErr_SetString( PyExc_MemoryError,
@@ -75,6 +66,16 @@ APPLY_SPECIFIC(dnn_sptf_grid)(PyGpuArrayObject * theta,
     height = (int) *( (npy_int64 *) PyArray_GETPTR1( out_dims, 2 ) );
     width = (int) *( (npy_int64 *) PyArray_GETPTR1( out_dims, 3 ) );
 
+    if ( PyGpuArray_DIM( theta, 0 ) != num_images ||
+         PyGpuArray_DIM( theta, 1 ) != 2 || PyGpuArray_DIM( theta, 2 ) != 3 )
+    {
+        PyErr_Format( PyExc_RuntimeError,
+            "GpuDnnTransformerGrid: incorrect dimensions for theta, expected (%d, %d, %d), got (%d, %d, %d)",
+            num_images, 2, 3, PyGpuArray_DIMS( theta )[0],
+            PyGpuArray_DIMS( theta )[1], PyGpuArray_DIMS( theta )[2] );
+        return 1;
+    }
+
     // Set transformed output dimensions to setup the descriptor
     desc_dims[0] = num_images;
     desc_dims[1] = num_channels;

theano/gpuarray/tests/test_dnn.py

@@ -3,6 +3,7 @@ import logging
 from collections import OrderedDict
 
 from nose.plugins.skip import SkipTest
+from nose.tools import assert_raises
 from parameterized import parameterized
 import numpy as np
 from itertools import product, chain
@@ -2470,6 +2471,34 @@ def test_dnn_spatialtf():
     utt.assert_allclose(img_out_cpu, img_out_gpu, atol=atol, rtol=rtol)
 
 
+def test_dnn_spatialtf_invalid_shapes():
+    if not dnn.dnn_available(test_ctx_name):
+        raise SkipTest(dnn.dnn_available.msg)
+
+    inputs = T.tensor4('inputs')
+    theta = T.tensor3('theta')
+
+    st_dnn = dnn.dnn_spatialtf(inputs, theta)
+    st_dnn_func = theano.function([inputs, theta], st_dnn)
+
+    inputs_val = np.ones((3, 5, 7, 7), dtype=theano.config.floatX)
+
+    def try_theta_shp(theta_shp):
+        theta_val = np.ones(theta_shp, dtype=theano.config.floatX)
+        return st_dnn_func(inputs_val, theta_val)
+
+    # the theta shape for this input should be (3, 2, 3)
+    try_theta_shp((3, 2, 3))
+
+    # incorrect parameter dimensions
+    assert_raises(RuntimeError, try_theta_shp, (3, 1, 3))
+    assert_raises(RuntimeError, try_theta_shp, (3, 2, 1))
+
+    # number of rows does not match the number of input rows
+    assert_raises(RuntimeError, try_theta_shp, (1, 2, 3))
+    assert_raises(RuntimeError, try_theta_shp, (4, 2, 3))
+
+
 def test_dnn_spatialtf_grad():
     if not dnn.dnn_available(test_ctx_name):
         raise SkipTest(dnn.dnn_available.msg)