Fit split axes same with transpose after optimizer (#1918)

Signed-off-by: Deyu Huang <[email protected]>

Author: hwangdeyu

tf2onnx/optimizer/transpose_optimizer.py

@@ -105,7 +105,7 @@ def _calculate_new_shape(graph, op):
                 new_shape = [input_shape[p] for p in perm]
                 return graph.make_const(utils.make_name("new_shape"), np.array(new_shape, dtype=np.int64)).output[0]
 
-            # reshape requires tha output shape can only contain one -1, if not some extra op needed.
+            # reshape requires the output shape can only contain one -1, if not some extra op needed.
             input_shape = graph.make_node("Shape", [op.input[0]]).output[0]
             indice = graph.make_const(utils.make_name("indice"), np.array(perm, np.int64)).output[0]
 
@@ -668,9 +668,12 @@ def _concat_handler(self, trans, node):
         return False
 
     def _split_handler(self, trans, node):
-        # Todo: need handle cases where Slit node has more than 1 outputs.
+        # Todo: need handle cases where Split node has more than 1 outputs.
         if self._handle_node_having_branches(trans, node):
-            node.set_attr("axis", 1)
+            perm = trans.get_attr_value("perm")
+            axis = node.get_attr_value("axis", 0)
+            new_axis = perm[axis]
+            node.set_attr("axis", new_axis)
             return True
         return False