Merge pull request #4691 from JesseLivezey/corrmm_dtype

theano/tensor/nnet/corr.py

@@ -85,6 +85,14 @@ class BaseCorrMM(gof.OpenMPOp):
             str(self.subsample),
             str(self.filter_dilation))
 
+    @staticmethod
+    def as_common_dtype(in1, in2):
+        """
+        Upcast input variables if neccesary.
+        """
+        dtype = theano.scalar.upcast(in1.dtype, in2.dtype)
+        return in1.astype(dtype), in2.astype(dtype)
+
     def c_support_code(self):
         ccodes = blas_headers.blas_header_text()
         if self.blas_type == 'openblas':
@@ -420,6 +428,7 @@ class CorrMM(BaseCorrMM):
     def make_node(self, img, kern):
         img = as_tensor_variable(img)
         kern = as_tensor_variable(kern)
+        img, kern = self.as_common_dtype(img, kern)
         if img.type.ndim != 4:
             raise TypeError('img must be 4D tensor')
         if kern.type.ndim != 4:
@@ -476,6 +485,7 @@ class CorrMM_gradWeights(BaseCorrMM):
     def make_node(self, img, topgrad, shape=None):
         img = as_tensor_variable(img)
         topgrad = as_tensor_variable(topgrad)
+        img, topgrad = self.as_common_dtype(img, topgrad)
         if img.type.ndim != 4:
             raise TypeError('img must be 4D tensor')
         if topgrad.type.ndim != 4:
@@ -572,6 +582,7 @@ class CorrMM_gradInputs(BaseCorrMM):
     def make_node(self, kern, topgrad, shape=None):
         kern = as_tensor_variable(kern)
         topgrad = as_tensor_variable(topgrad)
+        kern, topgrad = self.as_common_dtype(kern, topgrad)
         if kern.type.ndim != 4:
             raise TypeError('kern must be 4D tensor')
         if topgrad.type.ndim != 4:

theano/tensor/nnet/tests/test_corr.py

@@ -2,6 +2,7 @@ from __future__ import absolute_import, print_function, division
 
 from nose.plugins.skip import SkipTest
 from nose.plugins.attrib import attr
+from nose.tools import assert_equals
 import numpy
 from six import integer_types
 
@@ -259,6 +260,32 @@ class TestCorr2D(utt.InferShapeTester):
         self.assertRaises(Exception, self.validate, (3, 2, 8, 8), (4, 2, 5, 5),
                           'valid', input=T.dtensor3())
 
+    def test_dtype_upcast(self):
+        """
+        Checks dtype upcast for CorrMM methods.
+        """
+        def rand(shape, dtype='float64'):
+            r = numpy.asarray(numpy.random.rand(*shape), dtype=dtype)
+            return r * 2 - 1
+
+        ops = [corr.CorrMM, corr.CorrMM_gradWeights, corr.CorrMM_gradInputs]
+        a_shapes = [[4, 5, 6, 3], [1, 5, 6, 3], [1, 5, 6, 3]]
+        b_shapes = [[7, 5, 3, 2], [1, 5, 3, 1], [7, 1, 3, 1]]
+        dtypes = ['float32', 'float64']
+
+        for op, a_shape, b_shape in zip(ops, a_shapes, b_shapes):
+            for a_dtype in dtypes:
+                for b_dtype in dtypes:
+                    c_dtype = theano.scalar.upcast(a_dtype, b_dtype)
+                    a_tens = T.tensor4(dtype=a_dtype)
+                    b_tens = T.tensor4(dtype=b_dtype)
+                    a_tens_val = rand(a_shape, dtype=a_dtype)
+                    b_tens_val = rand(b_shape, dtype=b_dtype)
+
+                    c_tens = op()(a_tens, b_tens)
+                    f = theano.function([a_tens, b_tens], c_tens, mode=self.mode)
+                    assert_equals(f(a_tens_val, b_tens_val).dtype, c_dtype)
+
     @attr('slow')
     def test_infer_shape_forward(self):
         if theano.config.mode == "FAST_COMPILE":