Merge pull request #308 from zhijxu-MS/tmp_branch_for_PR2

add tf.eye supported partially

Author: nbcsm

tests/test_backend.py

@@ -39,6 +39,8 @@
 _OUTPUT = "output:0"
 _TFOUTPUT1 = "output1"
 _OUTPUT1 = "output1:0"
+_TFOUTPUT2 = "output2"
+_OUTPUT2 = "output2:0"
 
 
 def make_xval(shape):
@@ -127,6 +129,63 @@ def test_expand_dims_more_unknown_rank(self):
         _ = tf.identity(op, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(9, "ConstantOfShape")
+    def test_eye_non_const1(self):
+        # tf.eye(num_rows), num_rows is not const here
+        tf.reset_default_graph()
+        x_val = np.array(5, dtype=np.int32)
+        x = tf.placeholder(tf.int32, shape=[], name=_TFINPUT)
+        y = tf.eye(x, dtype=tf.int32)
+        _ = tf.identity(y, name=_TFOUTPUT)
+        y1 = tf.eye(x, dtype=tf.int64)
+        _ = tf.identity(y1, name=_TFOUTPUT1)
+        y2 = tf.eye(x, dtype=tf.float32)
+        _ = tf.identity(y2, name=_TFOUTPUT2)
+        self._run_test_case([_OUTPUT, _OUTPUT1, _OUTPUT2], {_INPUT: x_val}, rtol=0)
+
+        # tf.eye(num_rows, num_columns), both num_rows and num_columns are not const here
+        tf.reset_default_graph()
+        x_val = np.array([5, 10], dtype=np.int32)
+        x = tf.placeholder(tf.int32, shape=[2], name=_TFINPUT)
+        y = tf.eye(x[0], x[1], dtype=tf.int32)
+        _ = tf.identity(y, name=_TFOUTPUT)
+        y1 = tf.eye(x[0], x[1], dtype=tf.int64)
+        _ = tf.identity(y1, name=_TFOUTPUT1)
+        y2 = tf.eye(x[0], x[1], dtype=tf.float32)
+        _ = tf.identity(y2, name=_TFOUTPUT2)
+        self._run_test_case([_OUTPUT, _OUTPUT1, _OUTPUT2], {_INPUT: x_val}, rtol=0)
+
+    @check_tf_min_version("1.11", "eye has bug when version is below 1.11")
+    @check_opset_min_version(9, "ConstantOfShape")
+    def test_eye_non_const2(self):
+        # tf.eye(num_rows), num_rows is not const here
+        for np_dtype, tf_dtype in zip([np.int32, np.int64, np.float32, np.float64],
+                                      [tf.int32, tf.int64, tf.float32, tf.float64]):
+            tf.reset_default_graph()
+            x_val = np.array(5, dtype=np_dtype)
+            x = tf.placeholder(tf_dtype, shape=[], name=_TFINPUT)
+            y = tf.eye(x, dtype=tf.int32)
+            _ = tf.identity(y, name=_TFOUTPUT)
+            y1 = tf.eye(x, dtype=tf.int64)
+            _ = tf.identity(y1, name=_TFOUTPUT1)
+            y2 = tf.eye(x, dtype=tf.float32)
+            _ = tf.identity(y2, name=_TFOUTPUT2)
+            self._run_test_case([_OUTPUT, _OUTPUT1, _OUTPUT2], {_INPUT: x_val}, rtol=0)
+
+        # tf.eye(num_rows, num_columns), both num_rows and num_columns are not const here
+        for np_dtype, tf_dtype in zip([np.int32, np.int64, np.float32, np.float64],
+                                      [tf.int32, tf.int64, tf.float32, tf.float64]):
+            tf.reset_default_graph()
+            x_val = np.array([5, 10], dtype=np_dtype)
+            x = tf.placeholder(tf_dtype, shape=[2], name=_TFINPUT)
+            y = tf.eye(x[0], x[1], dtype=tf.int32)
+            _ = tf.identity(y, name=_TFOUTPUT)
+            y1 = tf.eye(x[0], x[1], dtype=tf.int64)
+            _ = tf.identity(y1, name=_TFOUTPUT1)
+            y2 = tf.eye(x[0], x[1], dtype=tf.float32)
+            _ = tf.identity(y2, name=_TFOUTPUT2)
+            self._run_test_case([_OUTPUT, _OUTPUT1, _OUTPUT2], {_INPUT: x_val}, rtol=0)
+
     @check_opset_min_version(7, "trig")
     def test_trig_ops(self):
         for op in [tf.sin, tf.cos, tf.tan, tf.asin, tf.acos, tf.atan]:

tf2onnx/rewriter/__init__.py

@@ -9,6 +9,7 @@
 from tf2onnx.rewriter.cond_rewriter import rewrite_cond
 from tf2onnx.rewriter.random_uniform import rewrite_random_uniform, rewrite_random_uniform_fold_const
 from tf2onnx.rewriter.leakyrelu_rewriter import rewrite_leakyrelu
+from tf2onnx.rewriter.eye_rewriter import rewrite_eye
 from tf2onnx.rewriter.thresholded_relu_rewriter import rewrite_thresholded_relu
 from tf2onnx.rewriter.rnn import rewrite_single_direction_lstm, rewrite_bi_direction_lstm, \
     rewrite_single_direction_gru, rewrite_bi_direction_gru, \
@@ -20,6 +21,7 @@
     "rewrite_random_uniform_fold_const",
     "rewrite_leakyrelu",
     "rewrite_thresholded_relu",
+    "rewrite_eye",
     "rewrite_single_direction_lstm",
     "rewrite_bi_direction_lstm",
     "rewrite_single_direction_gru",

tf2onnx/rewriter/eye_rewriter.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter.eye_rewriter - supports tf.eye
+"""
+
+from onnx import onnx_pb
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+# pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable
+
+
+def rewrite_eye(g, ops):
+    # schema of eye is eye(num_rows, num_columns=None), if num_columns not specified then it's equal to num_rows
+    # tf.eye is implemented by a sub_graph which contains op "MatrixDiag" or "MatrixSetDiag" while
+    # these two ops are un-supported directly in onnx
+    # but onnx op EyeLike can be used to map the sub_graph
+    # "rewrite_eye" supports tf.eye(non_const) and tf.eye(non_const1, non_const2).
+    # tf.eye(const) and tf.eye(const1, const2) are not supported in this rewriter
+
+    # ConstantOfShape in opset 9 is used, so if opset less than 9 then do nothing
+    if g.opset < 9:
+        return g.get_nodes()
+
+    pattern1 = \
+        OpTypePattern("MatrixDiag", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("Const", name="fill_value"),
+                OpTypePattern("ConcatV2", inputs=[
+                    "*",
+                    "*",
+                    OpTypePattern("Pack", inputs=[
+                        OpTypePattern("Minimum|Cast", name="min_or_cast")
+                    ])
+                ])
+            ])
+        ])
+    pattern2 = \
+        OpTypePattern("MatrixSetDiag", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill"),
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("Const", name="fill_value"),
+                OpTypePattern("ConcatV2", inputs=[
+                    "*",
+                    "*",
+                    OpTypePattern("Pack", inputs=[
+                        OpTypePattern("Minimum|Cast", name="min_or_cast")
+                    ])
+                ])
+            ])
+        ])
+
+    for pattern in [pattern1, pattern2]:
+        matcher = GraphMatcher(pattern, allow_reorder=True)
+        match_results = list(matcher.match_ops(ops))
+        for match_result in match_results:
+            if match_result.get_op("fill_value").get_tensor_value() != 1:
+                continue
+
+            min_or_cast = match_result.get_op("min_or_cast")
+            if min_or_cast.type == "Minimum":
+                min_node = min_or_cast
+            elif min_or_cast.type == "Cast" and min_or_cast.inputs[0].type == "Minimum":
+                min_node = min_or_cast.inputs[0]
+            else:
+                continue
+
+            num_rows = min_node.inputs[0]
+            num_columns = min_node.inputs[1]
+
+            old_output = match_result.get_op("output_eye_matrix")
+            output_dtypes = [g.get_dtype(old_output.output[0])]
+            output_shapes = [g.get_shape(old_output.output[0])]
+            g.remove_node(old_output.name)
+
+            # onnx op "EyeLike" need a 2D tensor, so generate it
+            num_rows = g.make_node("Unsqueeze", num_rows.output, attr={"axes": [0]})
+            num_columns = g.make_node("Unsqueeze", num_columns.output, attr={"axes": [0]})
+            matrix_shape = g.make_node("Concat", [num_rows.output[0], num_columns.output[0]], attr={"axis": 0})
+            # cast nodes added for "ConstantOfShape" in ONNX only accepts int64 data.
+            matrix_shape_int64 = g.make_node("Cast", matrix_shape.output, attr={"to": onnx_pb.TensorProto.INT64})
+            zero_matrix = g.make_node("ConstantOfShape", matrix_shape_int64.output)
+
+            new_output = g.make_node("EyeLike", zero_matrix.output, attr={"dtype": output_dtypes[0]},
+                                     name=old_output.name, shapes=output_shapes, dtypes=output_dtypes)
+
+    return g.get_nodes()

tf2onnx/tfonnx.py

@@ -771,7 +771,7 @@ def compat_handler(ctx, node, **kwargs):
     # bi-directional re-writer should be placed after single directional re-writer
     rewriters = [rewrite_transpose, rewrite_flatten,
                  rewrite_random_uniform, rewrite_random_uniform_fold_const,
-                 rewrite_random_normal, rewrite_dropout,
+                 rewrite_random_normal, rewrite_dropout, rewrite_eye,
                  rewrite_leakyrelu, rewrite_thresholded_relu, rewrite_conv2d_with_pad,
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,
                  rewrite_single_direction_gru, rewrite_bi_direction_gru,