refactor

move is_xx_op from rnn_utils to utils

Author: zhijxu-MS

tf2onnx/rewriter/bigru_rewriter.py

@@ -13,7 +13,7 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.rewriter.rnn_utils import is_reverse_op
+from tf2onnx.utils import is_reverse_op
 from tf2onnx.rewriter.bilstm_rewriter import slice_bilstm_for_original_lstm_consumers,\
      get_reverse_nodes_after_y_output, get_np_val_for_const, _process_single_init_node
 

tf2onnx/rewriter/bilstm_rewriter.py

@@ -13,7 +13,7 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.rewriter.rnn_utils import is_reverse_op
+from tf2onnx.utils import is_reverse_op
 from tf2onnx.graph_builder import GraphBuilder
 
 logger = logging.getLogger(__name__)

tf2onnx/rewriter/loop_rewriter_base.py

@@ -12,8 +12,8 @@
 from collections import OrderedDict
 from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
-from tf2onnx.rewriter.rnn_utils import is_loopcond_op, is_tensor_array_op
-from tf2onnx.rewriter.rnn_utils import is_tensor_array_gather_op, is_tensor_array_write_op
+from tf2onnx.utils import is_loopcond_op, is_tensor_array_op
+from tf2onnx.utils import is_tensor_array_gather_op, is_tensor_array_write_op
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT
 from tf2onnx.utils import TensorValueInfo
 

tf2onnx/rewriter/lstm_rewriter.py

@@ -13,8 +13,8 @@
 import numpy as np
 from tf2onnx import utils
 from tf2onnx.graph_builder import GraphBuilder
-from tf2onnx.rewriter.rnn_utils import RNNUnitType, RnnWeight, \
-    is_concat_op, is_slice_op, get_weights_from_const_node
+from tf2onnx.rewriter.rnn_utils import RNNUnitType, RnnWeight, get_weights_from_const_node
+from tf2onnx.utils import is_concat_op, is_slice_op
 
 from tf2onnx.rewriter.unit_rnn_rewriter_base import UnitRnnRewriterBase
 

tf2onnx/rewriter/rnn_utils.py

@@ -263,35 +263,3 @@ def get_weights_from_const_node(g, node):
         return None
 
     return RnnWeight(node, val, dtype)
-
-
-def is_reverse_op(op):
-    return op.type in ("ReverseV2", "ReverseSequence")
-
-
-def is_concat_op(op):
-    return op.type in ("Concat", "ConcatV2", "ConcatV3")
-
-
-def is_tensor_array_gather_op(op):
-    return op.type in ("TensorArrayGatherV2", "TensorArrayGatherV3")
-
-
-def is_tensor_array_write_op(op):
-    return op.type in ("TensorArrayWriteV2", "TensorArrayWriteV3")
-
-
-def is_tensor_array_op(op):
-    return op.type in ("TensorArrayV2", "TensorArrayV3")
-
-
-def is_loopcond_op(op):
-    return op.type == "LoopCond"
-
-
-def is_select_op(op):
-    return op.type == "Select"
-
-
-def is_slice_op(op):
-    return op.type == "Slice"

tf2onnx/rewriter/unit_rnn_rewriter_base.py

@@ -13,12 +13,11 @@
 from tf2onnx.graph_builder import GraphBuilder
 from tf2onnx.rewriter.loop_rewriter_base import LoopRewriterBase, Context
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT, get_pattern, \
-    get_rnn_scope_name, parse_rnn_loop, is_select_op, is_tensor_array_write_op, \
-    seq_len_pattern
+    get_rnn_scope_name, parse_rnn_loop, seq_len_pattern
+from tf2onnx.utils import is_select_op, is_tensor_array_write_op
 from tf2onnx.graph_matcher import GraphMatcher
 
 
-
 logger = logging.getLogger(__name__)
 
 

tf2onnx/shape_inference.py

@@ -12,7 +12,6 @@
 import numpy as np
 from onnx import onnx_pb
 from tf2onnx import utils
-from tf2onnx.rewriter import rnn_utils
 
 # pylint: disable=logging-not-lazy,missing-docstring,consider-swap-variables
 
@@ -215,7 +214,7 @@ def infer_input_shapes(g, node):
 def infer_output_shapes_with_partial_inputs(g, node):
     # output shape of concat op: only the dim val of concatenated dim will be changed
     # so only partial(at least one) input shapes need to be known to infer output shape of concat node
-    if rnn_utils.is_concat_op(node):
+    if utils.is_concat_op(node):
         data_inputs = node.input[:-1]
         input_shapes = [g.get_shape(node) for node in data_inputs]
         input_shapes = [shape for shape in input_shapes if shape is not None]

tf2onnx/utils.py

@@ -493,3 +493,35 @@ def get_url(url, path, max_retries=5):
 
     with open(path, "wb") as f:
         f.write(response.content)
+
+
+def is_reverse_op(op):
+    return op.type in ("ReverseV2", "ReverseSequence")
+
+
+def is_concat_op(op):
+    return op.type in ("Concat", "ConcatV2", "ConcatV3")
+
+
+def is_tensor_array_gather_op(op):
+    return op.type in ("TensorArrayGatherV2", "TensorArrayGatherV3")
+
+
+def is_tensor_array_write_op(op):
+    return op.type in ("TensorArrayWriteV2", "TensorArrayWriteV3")
+
+
+def is_tensor_array_op(op):
+    return op.type in ("TensorArrayV2", "TensorArrayV3")
+
+
+def is_loopcond_op(op):
+    return op.type == "LoopCond"
+
+
+def is_select_op(op):
+    return op.type == "Select"
+
+
+def is_slice_op(op):
+    return op.type == "Slice"