dimshuffle on a broadcastable dimension is handled

theano/tensor/opt.py

@@ -606,8 +606,22 @@ def local_dimshuffle_lift(node):
             return [ret]
 
     # remove useless dimshuffle in general
-    if (list(op.new_order) == list(range(len(op.new_order))) and
-            len(op.new_order) == input.type.ndim):
+    # covers two types of useless dimshuffle:
+    #   1 - dimshuffle all dimensions in order
+    #   2 - dimshuffle a broadcastable dimension
+    is_useless = True
+    all_broadcastable_dims = [i for (i, is_broadcastable)
+                              in enumerate(input.type.broadcastable)
+                              if is_broadcastable] + ['x']
+    for i in range(input.type.ndim):
+        if (op.new_order[i] == i or
+                (i in all_broadcastable_dims and
+                 op.new_order[i] in all_broadcastable_dims)):
+            continue
+        else:
+            is_useless = False
+            break
+    if is_useless:
         return [input]
 
 

theano/tensor/tests/test_opt.py

@@ -202,6 +202,18 @@ class test_dimshuffle_lift(unittest.TestCase):
         # Check stacktrace was copied over correctly after opt was applied
         self.assertTrue(hasattr(g.outputs[0].tag, 'trace'))
 
+    def test_dimshuffle_on_broadcastable(self):
+        x, y, z = inputs([False, True], [True, False, True], [False, False, True])
+        ds_x = ds(x, (0, 'x'))   # useless
+        ds_y = ds(y, (2, 1, 0))  # useless
+        ds_z = ds(z, (2, 1, 0))  # usefull
+        g = FunctionGraph([x, y, z], [ds_x, ds_y, ds_z])
+        self.assertTrue(str(g) == "[DimShuffle{0,x}(x), DimShuffle{2,1,0}(y), DimShuffle{2,1,0}(z)]")
+        dimshuffle_lift.optimize(g)
+        self.assertTrue(str(g) == "[x, y, DimShuffle{2,1,0}(z)]")
+        # Check stacktrace was copied over correctly after opt was applied
+        self.assertTrue(hasattr(g.outputs[0].tag, 'trace'))
+
 
 def test_add_canonizer_problem0():
     n_segments = 10