Merge pull request #3 from nouiz/lamblin-fix_1507

theano/tensor/blas.py

@@ -1906,59 +1906,53 @@ def local_dot22_to_dot22scalar(node):
     d = node.inputs[dot22_idx]
     i_scalar = [_as_scalar(x, dtype=d.dtype) for x in node.inputs]
     if not any(i_scalar):
-        i_mul = [x.owner and x.owner.op == T.mul for x in node.inputs]
+        # Check if we can reorder the graph as this mul have a mul in inputs.
+        # We support only 1 additional level of mul.
+        # The canonizer should have merged those mul together.
+        i_mul = [x.owner and x.owner.op == T.mul and
+                 any([_as_scalar(x_i, dtype=d.dtype)
+                   for x_i in x.owner.inputs])
+                 for x in node.inputs]
         if not any(i_mul):
             #no scalar in input and no multiplication
             #if their was a multiplication we couls reorder the graph
             #by the associativity of the graph.
             return False
 
-        #maybe we can reorder the graph as this mul have a mul in input.
-        #The canonizer should have merged those mul together.
-        #We support only 1 additional level of mul.
-        mul_idx = i_mul.index(True)  # we take the first mul!
+        mul_idx = i_mul.index(True)  # The first one should always work
         m = node.inputs[mul_idx]
 
-        if any([_as_scalar(x, dtype=d.dtype)
-                for x in m.owner.inputs]):
-            scalar_idx = -1
-            for i, x in enumerate(m.owner.inputs):
-                if _as_scalar(x, dtype=d.dtype) 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],
-                                  dtype=d.dtype), d.type.dtype)
-            assert not a.type.ndim
-            dot = _dot22scalar(d.owner.inputs[0], d.owner.inputs[1], a)
-
-            # The other inputs to the original node that were
-            # neither part of the dot22 or this mul should be
-            # factors in the returned "mul" node.
-            assert dot22_idx != mul_idx
-            other_factors = [inpt
-                             for i, inpt in enumerate(node.inputs)
-                             if i not in (dot22_idx, mul_idx)]
-            other_m_inputs = [inpt
-                              for i, inpt in enumerate(m.owner.inputs)
-                              if i != scalar_idx]
-
-            return [T.mul(dot, *(other_factors + other_m_inputs))]
-        elif m.owner and m.owner.op == T.mul:
-            _logger.info('Not optimizing dot22 with inputs %s %s %s %s. '
-                    'we need to check in a recursive way in the mul if we can '
-                    'reorder the graph. The canonizer should have done this.',
-                    d, m, d.type, m.type)
-        else:
+        scalar_idx = -1
+        for i, x in enumerate(m.owner.inputs):
+            if _as_scalar(x, dtype=d.dtype) 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],
+                              dtype=d.dtype), d.type.dtype)
+        assert not a.type.ndim
+        dot = _dot22scalar(d.owner.inputs[0], d.owner.inputs[1], a)
+
+        # The other inputs to the original node that were
+        # neither part of the dot22 or this mul should be
+        # factors in the returned "mul" node.
+        assert dot22_idx != mul_idx
+        other_factors = [inpt
+                         for i, inpt in enumerate(node.inputs)
+                         if i not in (dot22_idx, mul_idx)]
+        other_m_inputs = [inpt
+                          for i, inpt in enumerate(m.owner.inputs)
+                          if i != scalar_idx]
+
+        return [T.mul(dot, *(other_factors + other_m_inputs))]
 
     scalar_idx = -1
     for i, x in enumerate(node.inputs):

theano/tensor/tests/test_blas.py

@@ -1024,6 +1024,60 @@ def test_dot22scalar_cast():
             assert _dot22scalar in [x.op for x in f.maker.fgraph.toposort()]
 
 
+def test_local_dot22_to_dot22scalar():
+    """
+    This test that the bug in gh-1507 is really fixed
+    """
+    A = T.dmatrix()
+    mode = theano.compile.mode.get_default_mode()
+    opt = theano.tensor.opt.in2out(
+        theano.tensor.blas.local_dot22_to_dot22scalar)
+    mode = mode.__class__(optimizer=opt)
+
+    x = T.dscalar()
+    y = T.dscalar()
+    z = T.dscalar()
+    # make sure to don't have dimshuffle as we don't opt those cases
+    m = T.dmatrix()
+    r = T.drow()
+    for idx, node in enumerate([
+        #Old working cases
+        T.mul(_dot22(A, A), x),
+        T.mul(_dot22(A, A), x, y),
+        T.mul(_dot22(A, A), x, r),
+        T.mul(_dot22(A, A), m, x),
+        T.mul(_dot22(A, A), x, m),
+        T.mul(_dot22(A, A), x, (m * y)),
+        T.mul(_dot22(A, A), (m * y), x),
+        T.mul(_dot22(A, A), x, (r * y)),
+        T.mul(_dot22(A, A), (r * y), x),
+        T.mul(_dot22(A, A), (x * y), (m * x)),
+        T.mul(_dot22(A, A), (r * y), (y * x)),
+
+        # Case that was raising an assert that is fixed in gh-1507
+        T.mul(_dot22(A, A), (m * y), m),
+        T.mul(_dot22(A, A), m, (m * y)),
+        T.mul(_dot22(A, A), (r * y), (m * x)),
+
+        # assert fixed in gh-1507 and opt case added in gh-1515
+        T.mul(_dot22(A, A), (m * y * z), m),
+        T.mul(_dot22(A, A), m, (m * y * z)),
+
+        # Opt case added in gh-1515
+        T.mul(_dot22(A, A), T.mul(m, y, z), m),
+        T.mul(_dot22(A, A), m, T.mul(m, y, z)),
+
+        #Case that opt later in gh-1515
+        T.mul(_dot22(A, A), (r * m), (m * x)),
+    ]):
+        node2 = theano.tensor.blas.local_dot22_to_dot22scalar.transform(
+            node.owner)
+        assert node2
+        f = theano.function([x, y, z, m, r, A], node,
+                            mode=mode, on_unused_input='ignore')
+        f(.1, .2, .3, [[1, 2], [3, 4]], [[5, 6]], [[7, 8], [9, 10]])
+
+
 def test_dot_w_self():
     # This can trigger problems in the optimization because what would
     # normally be a gemm must not be because the output is aliased to