Make dot22scalar work in python with complex.

theano/tensor/blas.py

@@ -1503,6 +1503,9 @@ class Dot22Scalar(GemmRelated):
     def c_code(self, node, name, inp, out, sub): #DEBUG
         _x, _y, _a = inp
         _zout, = out
+        if node.inputs[0].type.dtype.startswith('complex'):
+            raise utils.MethodNotDefined('%s.c_code' \
+                    % self.__class__.__name__)
         if len(self.c_libraries())<=0:
             return super(Dot22Scalar, self).c_code(node, name, (_x, _y), (_zout, ), sub)
         full_code = self.build_gemm_call() % dict(locals(), **sub)
@@ -1551,12 +1554,18 @@ def local_dot22_to_dot22scalar(node):
         m = node.inputs[mul_idx]
 
         if len(m.owner.inputs)==2 and any([_as_scalar(x) for x in m.owner.inputs]):
-            scalar_idx = 0
+            scalar_idx = -1
             for i,x in enumerate(m.owner.inputs):
-                if _as_scalar(x):
+                if _as_scalar(x) and (theano.scalar.upcast(x.type.dtype,d.type.dtype)
+                                      == d.type.dtype):
                     scalar_idx=i
                     break
 
+            if scalar_idx<0:
+                _logger.info('Not optimizing dot22 with inputs %s %s, as the type '
+                             'of the scalar cannot be upcasted to the matrix type',
+                             node.inputs, [x.type for x in node.inputs])
+                return False
             a = T.cast(_as_scalar(m.owner.inputs[scalar_idx]), d.type.dtype)
             assert not a.type.ndim
             dot=_dot22scalar(d.owner.inputs[0], d.owner.inputs[1], a)

theano/tensor/tests/test_blas.py

@@ -704,53 +704,81 @@ def test_dot22scalar():
     ## TODO: exclude other optimizations in BlasOpt?
     #m = theano.compile.get_default_mode().including('local_dot_to_dot22','local_dot22_to_dot22scalar','specialize')
     #m = theano.compile.get_default_mode().including('BlasOpt', 'specialize')
-    a=T.matrix()
-    b=T.matrix()
-    c=T.matrix()
     rng = numpy.random.RandomState(unittest_tools.fetch_seed())
+    for dtype1 in ['complex64', 'complex128']:
+        a=T.matrix('a', dtype = dtype1)
+        for dtype2 in ['complex64', 'complex128']:
+            b=T.matrix('b', dtype = dtype2)
+            for dtype3 in ['complex64', 'complex128']:
+                c=T.matrix('c', dtype = dtype3)
+                for dtype4 in ['complex64', 'complex128']:
+                    cst = theano.tensor.basic.constant(.2, dtype=dtype4)
+                    cst2 = theano.tensor.basic.constant(.1, dtype=dtype4)
+
+                    def check_dot22scalar(func, len_topo_scalar=-1):
+                        topo = func.maker.env.toposort()
+                        ops = [x.op for x in topo]
+                        dtype4_upcast = theano.scalar.upcast(dtype4, dtype1, dtype2)
+                        if dtype1 == dtype2 == dtype3 == dtype4_upcast:
+                            if len_topo_scalar>0:
+                                assert len(topo) == len_topo_scalar
+                            assert _dot22scalar in ops, (dtype1, dtype2, dtype3, dtype4)
+                        elif dtype1 == dtype2 == dtype4_upcast:
+                            if not (len_topo_scalar > 0):
+                                assert len(topo) == len_topo_scalar
+                                assert _dot22scalar in ops, (dtype1, dtype2, dtype3, dtype4)
+                            else:
+                                # Currently there is a problem of optimization order
+                                # The constant get upcasted to float64 before we try to merge it
+                                # with the dot22 of float32. So this prevent the merge.
+                                assert _dot22scalar in ops or _dot22 in ops, (dtype1, dtype2, dtype3, dtype4)
+
+                        elif dtype1 == dtype2:
+                            assert _dot22 in ops, (dtype1, dtype2, dtype3, dtype4)
+                        else:
+                            assert T.dot in ops, (dtype1, dtype2, dtype3, dtype4)
+
 
                     def cmp(a_shp, b_shp, c_shp, sqr_shp=(5,5)):
-        av=rng.uniform(size=a_shp).astype(config.floatX)
-        bv=rng.uniform(size=b_shp).astype(config.floatX)
-        cv=rng.uniform(size=c_shp).astype(config.floatX)
-        sv=rng.uniform(size=sqr_shp).astype(config.floatX)
+                        av=rng.uniform(size=a_shp).astype(dtype1)
+                        bv=rng.uniform(size=b_shp).astype(dtype2)
+                        cv=rng.uniform(size=c_shp).astype(dtype3)
+                        sv=rng.uniform(size=sqr_shp).astype(dtype1)
 
-        if True:
-            f = theano.function([a,b],0.2*T.dot(a,b),mode=mode_blas_opt)
+                        if False:
+                            f = theano.function([a,b],cst*T.dot(a,b),mode=mode_blas_opt)
                             topo = f.maker.env.toposort()
-            assert _dot22scalar in [x.op for x in topo]
-            assert len(topo)==1
+                            check_dot22scalar(f, 1)
+
                             f(av,bv)
 
                         if True:
-            f = theano.function([a,b,c],0.2*c*T.dot(a,b),mode=mode_blas_opt)
+                            f = theano.function([a,b,c],cst*c*T.dot(a,b),mode=mode_blas_opt)
                             topo = f.maker.env.toposort()
-            assert _dot22scalar in [x.op for x in topo]
-            assert len(topo)==2
+                            check_dot22scalar(f, 2)
+
                             f(av,bv,cv)
 
-        f = theano.function([a,b,c],c * 0.2*T.dot(a,b),mode=mode_blas_opt)
+                        f = theano.function([a,b,c],c * cst*T.dot(a,b),mode=mode_blas_opt)
                         topo = f.maker.env.toposort()
-        assert _dot22scalar in [x.op for x in topo]
-        assert len(topo)==2
+                        check_dot22scalar(f, 2)
                         f(av,bv,cv)
 
                         ## Here, canonicalize also seems needed
                         ## TODO: add only the optimizations needed?
                         m2 = mode_blas_opt.including('canonicalize')
-        f = theano.function([a,b,c],0.1*c * 0.2*T.dot(a,b),mode=m2)
+                        f = theano.function([a,b,c],cst2 *c * cst*T.dot(a,b),mode=m2)
                         topo = f.maker.env.toposort()
-        assert _dot22scalar in [x.op for x in topo]
-        assert len(topo)==2
+                        check_dot22scalar(f, 2)
                         f(av,bv,cv)
 
-        f = theano.function([a,b,c],c * 0.2*a*T.dot(a,b),mode=m2)
+                        if dtype1 == dtype2 == dtype3:
+                            f = theano.function([a,b,c],c * cst*a*T.dot(a,b),mode=m2)
                             topo = f.maker.env.toposort()
-        assert _dot22scalar in [x.op for x in topo]
-        assert len(topo)==2
+                            check_dot22scalar(f, 2)
                             f(sv,sv,sv)
 
-        f = theano.function([a,b,c],0.2*c *a*T.dot(a,b),mode=mode_blas_opt)
+                            f = theano.function([a,b,c],cst*c *a*T.dot(a,b),mode=mode_blas_opt)
                             topo = f.maker.env.toposort()
                             #currently the canonizer don't always merge all Mul together...
                             # dot22scalar optimizer does not do a recursive search
@@ -761,10 +789,9 @@ def test_dot22scalar():
                             #    assert len(topo)==2
                             f(sv,sv,sv)
 
-        f = theano.function([a,b,c],c * a*0.2*T.dot(a,b),mode=m2)
+                            f = theano.function([a,b,c],c * a*cst*T.dot(a,b),mode=m2)
                             topo = f.maker.env.toposort()
-        assert _dot22scalar in [x.op for x in topo]
-        assert len(topo)==2
+                            check_dot22scalar(f, 2)
                             f(sv,sv,sv)
 
                     cmp((3,4),(4,5),(3,5))