add shape utilities and test

Author: nbcsm

tests/test_internals.py

@@ -17,7 +17,7 @@
 
 import tensorflow as tf
 import tf2onnx
-import tf2onnx.utils
+from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.graph import GraphUtil
 from common import unittest_main
@@ -50,14 +50,13 @@ def onnx_pretty(g, args=None):
 
 
 class Tf2OnnxInternalTests(unittest.TestCase):
-
     def setUp(self):
         """Setup test."""
         # suppress log info of tensorflow so that result of test can be seen much easier
         os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
         tf.logging.set_verbosity(tf.logging.WARN)
 
-        tf2onnx.utils.INTERNAL_NAME = 1
+        utils.INTERNAL_NAME = 1
         arg = namedtuple("Arg", "input inputs outputs verbose")
         self._args0 = arg(input="test", inputs=[], outputs=["output:0"], verbose=False)
         self._args1 = arg(input="test", inputs=["input:0"], outputs=["output:0"], verbose=False)
@@ -142,8 +141,8 @@ def test_rewrite_subgraph(self):
         for match in match_results:
             input_node = match.get_op('input')
             output_node = match.get_op('output')
-            op_name = tf2onnx.utils.make_name("ReplacedOp")
-            out_name = tf2onnx.utils.port_name(op_name)
+            op_name = utils.make_name("ReplacedOp")
+            out_name = utils.port_name(op_name)
             new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name)
             ops = g.replace_subgraph(ops, match, [], [output_node], [], [new_node])
         g.topological_sort(ops)
@@ -183,10 +182,26 @@ def test_cmdarg_parse(self):
         arg = "input/V-1_2:0,input/X:0[1,2,3],Y:1[4,5],Z:3,A:1,B"
         expected_inputs = ['input/V-1_2:0', 'input/X:0', 'Y:1', 'Z:3', 'A:1', 'B']
         expected_shape = {'Y:1': [4, 5], 'input/X:0': [1, 2, 3]}
-        inputs, shape_override = tf2onnx.utils.split_nodename_and_shape(arg)
+        inputs, shape_override = utils.split_nodename_and_shape(arg)
         self.assertEqual(expected_inputs, inputs)
         self.assertEqual(expected_shape, shape_override)
 
+    def test_shape_utils(self):
+        self.assertEqual(utils.merge_shapes(None, None), None)
+        self.assertEqual(utils.merge_shapes([], None), [])
+        self.assertEqual(utils.merge_shapes(None, [1, 2, 3]), [1, 2, 3])
+        self.assertEqual(utils.merge_shapes([1, 3], [None, 3]), [1, 3])
+        self.assertEqual(utils.merge_shapes([1, None, 3], (-1, 2, "unk")), [1, 2, 3])
+
+        self.assertTrue(utils.are_shapes_compatible(None, []))
+        self.assertTrue(utils.are_shapes_compatible([1, None, 3], (-1, 2, "unk")))
+        self.assertFalse(utils.are_shapes_compatible([1, 2, 3], (2, 3)))
+        self.assertFalse(utils.are_shapes_compatible([1, 2, 3], (4, 5, 6)))
+
+        self.assertTrue(utils.are_shapes_equal(None, None))
+        self.assertFalse(utils.are_shapes_equal(None, []))
+        self.assertTrue(utils.are_shapes_equal([1, 2, 3], (1, 2, 3)))
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/utils.py

@@ -342,11 +342,6 @@ def tf_name_scope(name):
     return '/'.join(name.split('/')[:-1])
 
 
-def create_vague_shape_like(shape):
-    make_sure(len(shape) >= 0, "rank should be >= 0")
-    return [-1 for i in enumerate(shape)]
-
-
 def get_temp_directory():
     return os.environ.get("TF2ONNX_TEMP_DIRECTORY", tempfile.mkdtemp())
 
@@ -364,3 +359,70 @@ def save_protobuf(path, message, as_text=False):
     else:
         with open(path, "wb") as f:
             f.write(message.SerializeToString())
+
+
+def is_list_or_tuple(obj):
+    return type(obj) in [list, tuple]
+
+
+def is_unknown_dimension(dim):
+    """  Return true if dim is not a positive integer value. """
+    if dim is None or not isinstance(dim, int):
+        return True
+    return dim <= 0
+
+
+def merge_shapes(shape1, shape2):
+    """
+    Merge 2 shapes, return merged shape, choose more specific dimension value from either side.
+    Raise exception for mismatch.
+    """
+    if shape1 is None:
+        return shape2
+    if shape2 is None:
+        return shape1
+
+    make_sure(is_list_or_tuple(shape1), "invalid type for shape1")
+    make_sure(is_list_or_tuple(shape2), "invalid type for shape2")
+    make_sure(len(shape1) == len(shape2), "shapes rank mismatch: shape1=%s, shape2=%s", shape1, shape2)
+
+    merged = []
+    for d1, d2 in zip(shape1, shape2):
+        d = d1
+        if is_unknown_dimension(d1):
+            d = d2
+        elif not is_unknown_dimension(d2):
+            make_sure(d1 == d2, "shapes dimension mismatch: shape1=%s, shape2=%s", shape1, shape2)
+        merged.append(d)
+    return merged
+
+
+def are_shapes_compatible(src, dest):
+    """
+    Returns True iff src is compatible with dest.
+    None is compatible with all shapes, different ranks are not considered as compatible
+    """
+    try:
+        merge_shapes(src, dest)
+        return True
+    except Exception:
+        return False
+
+
+def are_shapes_equal(src, dest):
+    if src is None:
+        return dest is None
+    if dest is None:
+        return src is None
+
+    make_sure(is_list_or_tuple(src), "invalid type for src")
+    make_sure(is_list_or_tuple(dest), "invalid type for dest")
+
+    if len(src) != len(dest):
+        return False
+    return all(i == j for i, j in zip(src, dest))
+
+
+def create_vague_shape_like(shape):
+    make_sure(len(shape) >= 0, "rank should be >= 0")
+    return [-1 for i in enumerate(shape)]