Merge pull request #978 from jignparm/jignparm/transpose_mul

Enhance Transpose optimizer "Mul" handler

Author: jignparm

tests/test_optimizers.py

@@ -691,7 +691,7 @@ def test_trans_can_be_replaced_with_reshape2(self):
                                        model_proto, remaining_transpose_num=0)
 
     def test_two_transposes_switch_with_mul(self):
-        const_node = self._make_onnx_const(np.array(10, dtype=np.float32), "const_10")
+        const_node = self._make_onnx_const(np.array(np.random.random(6), dtype=np.float32), "const_10")
         node0 = helper.make_node("Transpose", ["u1"], ["v1"], perm=[0, 2, 3, 1], name="trans_0")
         node1 = helper.make_node("Transpose", ["u2"], ["v2"], perm=[0, 2, 3, 1], name="trans_1")
 

tf2onnx/optimizer/transpose_optimizer.py

@@ -482,8 +482,21 @@ def _mul_handler(self, trans, node):
                 self._g.remove_node(node.name)
                 return True
 
-        # if the shape is () or (1), we just move transpose after the mul
-        if not multiplier.shape or (len(multiplier.shape) == 1 and multiplier.shape[0] == 1):
+        # if the shape is (), we just move transpose after the mul
+        if not multiplier.shape:
+            return self._switch_transpose_and_node(node, trans)
+
+        # if multiplier is 1-D
+        if len(multiplier.shape) == 1:
+            if multiplier.shape[0] == 1:
+                # shape is (1)
+                return self._switch_transpose_and_node(node, trans)
+
+            # shape is (N). reshape so that trans(shape) = 1,1,...,N
+            perm = list(trans.get_attr('perm').ints)
+            new_shape = np.ones(len(perm), dtype=np.int32)
+            new_shape[perm[-1]] = multiplier.shape[0]
+            multiplier_input_node.set_tensor_value(multiplier.reshape(new_shape))
             return self._switch_transpose_and_node(node, trans)
 
         return False