support MatrixDiag and MatrixSetDiag when they are used in tf.eye

Author: zhijxu-MS

tests/test_backend.py

@@ -127,6 +127,32 @@ 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(self):
+        # tf.eye(tf.shape)
+        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([[1.0, 2.0, -3.0, -4.0, 5.0]] * 2, dtype=np_dtype)
+            x = tf.placeholder(tf_dtype, shape=[None] * 2, name=_TFINPUT)
+            y_ = tf.eye(tf.shape(x)[0], dtype=tf.float32)
+            _ = tf.identity(y_, name=_TFOUTPUT)
+            y1_ = tf.eye(tf.shape(x)[1], dtype=tf.int32)
+            _ = tf.identity(y1_, name=_TFOUTPUT1)
+            self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: x_val}, rtol=0)
+
+        # tf.eye(tf.shape, tf.shape)
+        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([[1.0, 2.0, -3.0, -4.0, 5.0]] * 2, dtype=np_dtype)
+            x = tf.placeholder(tf_dtype, shape=[None] * 2, name=_TFINPUT)
+            y_ = tf.eye(tf.shape(x)[0], tf.shape(x)[1], dtype=tf.float32)
+            _ = tf.identity(y_, name=_TFOUTPUT)
+            y1_ = tf.eye(tf.shape(x)[0], tf.shape(x)[1], dtype=tf.int32)
+            _ = tf.identity(y1_, name=_TFOUTPUT1)
+            self._run_test_case([_OUTPUT, _OUTPUT1], {_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,76 @@
+# 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):
+    # 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(tf.shape(x)[i]) and tf.eye(tf.shape(x)[i], tf.shape(x)[j]).
+
+    # 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"),
+                OpTypePattern("ConcatV2", inputs=[
+                    "*",
+                    "*",
+                    OpTypePattern("Pack", inputs=[
+                        OpTypePattern("Minimum", name="min_node")
+                    ])
+                ])
+            ])
+        ])
+    pattern2 = \
+        OpTypePattern("MatrixSetDiag", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill"),
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("Const"),
+                OpTypePattern("ConcatV2", inputs=[
+                    "*",
+                    "*",
+                    OpTypePattern("Pack", inputs=[
+                        OpTypePattern("Minimum", name="min_node")
+                    ])
+                ])
+            ])
+        ])
+
+    for pattern in [pattern1, pattern2]:
+        matcher = GraphMatcher(pattern, allow_reorder=True)
+        match_results = list(matcher.match_ops(ops))
+        for match_result in match_results:
+            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)
+
+            min_node = match_result.get_op("min_node")
+            num_rows = min_node.inputs[0]
+            num_columns = min_node.inputs[1]
+
+            # 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,