Merge pull request #955 from RandySheriffH/rashuai/tf22

Enable tf22 ci

Author: RandySheriffH

ci_build/azure_pipelines/unit_test.yml

@@ -6,7 +6,7 @@ stages:
     - template: 'templates/job_generator.yml'
       parameters:
         python_versions: ['3.7']
-        tf_versions: ['1.14.0','1.15.2','2.1.0']
+        tf_versions: ['1.14.0','1.15.2','2.1.0','2.2.0']
         onnx_opsets: ['']
         job:
           steps:

tests/test_backend.py

@@ -1384,11 +1384,13 @@ def func():
         self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
 
     @skip_caffe2_backend()
+    @check_opset_after_tf_version("2.2", 9, "RandomUniform")
     def test_randomuniform_dyn_shape(self):
         # test for dynamic shape coming from a shape op
         x_val = np.array([0, 1, 2, 3, 5], dtype=np.int64)
         def func(x):
-            return random_uniform(x[3:], name=_TFOUTPUT, dtype=tf.float32)
+            ret = random_uniform(x[3:], dtype=tf.float32)
+            return tf.identity(ret, name=_TFOUTPUT)
         # since results are random, compare the shapes only
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, check_value=False, check_shape=True)
 
@@ -3338,7 +3340,7 @@ def test_matrix_set_diag_v3(self):
                                [7, 7, 7, 7]]]).astype(np.int64)
         diag_val = np.array([[1, 2, 3],
                              [4, 5, 6]]).astype(np.int64)
-        k_val = np.array([0])
+        k_val = np.array([0]).astype(np.int32)
 
         def func(base_matrix, diag, k):
             return tf.raw_ops.MatrixSetDiagV3(input=base_matrix, diagonal=diag, k=k, align='RIGHT_LEFT', name=_TFOUTPUT)

tf2onnx/onnx_opset/tensor.py

@@ -1835,8 +1835,16 @@ class MatrixDiagPartV2V3:
     @classmethod
     def version_11(cls, ctx, node, **kwargs):
         # assemble MatrixDiagPart V2&V3 by looping k diagonals with proper pads
+        const_zero = ctx.make_const(utils.make_name(node.name) + 'const_zero', np.array([0]).astype(np.int64))
+        const_one = ctx.make_const(utils.make_name(node.name) + 'const_one', np.array([1]).astype(np.int64))
+        const_two = ctx.make_const(utils.make_name(node.name) + 'const_two', np.array([2]).astype(np.int64))
+        const_neg_one = ctx.make_const(utils.make_name(node.name) + 'const_neg_one', np.array([-1]).astype(np.int64))
+        const_neg_two = ctx.make_const(utils.make_name(node.name) + 'const_neg_two', np.array([-2]).astype(np.int64))
+        def normalize():
+            raw_k = ctx.make_node('Cast', [node.input[1]], attr={'to': TensorProto.INT64}).output[0]
+            return ctx.make_node('Reshape', [raw_k, const_neg_one.output[0]]).output[0]
         input_tensor = node.input[0]
-        k = ctx.make_node('Cast', [node.input[1]], attr={'to': TensorProto.INT64}).output[0]
+        k = normalize()
         padding = node.input[2]
         align = 'LEFT_LEFT'
         if node.op.op_type == 'MatrixDiagPartV3':
@@ -1850,12 +1858,7 @@ def version_11(cls, ctx, node, **kwargs):
         for out in ctx.find_output_consumers(node.output[0]):
             if out.op.op_type == 'Identity':
                 ctx.set_shape(out.output[0], raw_output_shape)
-        # define constants
-        const_zero = ctx.make_const(utils.make_name(node.name) + 'const_zero', np.array([0]).astype(np.int64))
-        const_one = ctx.make_const(utils.make_name(node.name) + 'const_one', np.array([1]).astype(np.int64))
-        const_two = ctx.make_const(utils.make_name(node.name) + 'const_two', np.array([2]).astype(np.int64))
-        const_neg_one = ctx.make_const(utils.make_name(node.name) + 'const_neg_one', np.array([-1]).astype(np.int64))
-        const_neg_two = ctx.make_const(utils.make_name(node.name) + 'const_neg_two', np.array([-2]).astype(np.int64))
+
         # prepare new_shape of input
         input_shape = ctx.make_node('Shape', [input_tensor])
         shape_input_shape = ctx.make_node('Shape', [input_shape.output[0]])
@@ -2075,7 +2078,7 @@ def mkconsts(values, dtype=np.int64):
         xalign, yalign = align.split('_')
 
         # consts
-        const_zero_float, const_neg_one_float = mkconsts([[0], [-1]], np.float32)
+        const_zero_float, const_neg_one_float = mkconsts([0, -1], np.float32)
         const_zero, const_one, const_neg_one, const_neg_two, const_pad_vals, const_t = \
             mkconsts([[0], [1], [-1], [-2], pads, [-1, 1]])
         const_zero_scalar, const_one_scalar, const_neg_one_scalar = mkconsts([0, 1, -1])
@@ -2098,8 +2101,12 @@ def mkconsts(values, dtype=np.int64):
         m2_shape = ctx.make_node('Concat', [partial_shape, const_neg_one], attr={'axis': 0}).output[0]
         gather_shape = ctx.make_node('Concat', [partial_shape, const_one], attr={'axis': 0}).output[0]
 
+        def normalize():
+            raw_input1 = ctx.make_node('Cast', [node.input[1]], attr={'to': TensorProto.INT64}).output[0]
+            return ctx.make_node('Reshape', [raw_input1, const_neg_one])
+
         # get k0, k1 values. diags to be extracted
-        input1 = ctx.make_node('Cast', [node.input[1]], attr={'to': TensorProto.INT64})
+        input1 = normalize()
         k0 = ctx.make_node('ReduceMin', [input1.output[0]]).output[0]
         k1 = ctx.make_node('ReduceMax', [input1.output[0]]).output[0]
         k0_scalar = ctx.make_node('Squeeze', [k0]).output[0]

tf2onnx/rewriter/eye_rewriter.py

@@ -79,8 +79,38 @@ def rewrite_eye(g, ops):
                 OpTypePattern("Const", name="fill_value"),
             ]),
         ])
+    pattern5 = \
+        OpTypePattern("MatrixDiagV3", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("ConcatV2", inputs=[
+                    "*",
+                    OpTypePattern("ExpandDims", inputs=[
+                        OpTypePattern("Minimum|Cast", name="min_or_cast"),
+                        "*"
+                    ]),
+                    "*",
+                ]),
+                OpTypePattern("Const", name="fill_value"),
+            ]),
+            "*", "*", "*", "*",
+        ])
+    pattern6 = \
+        OpTypePattern("MatrixSetDiagV3", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill"),
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("ConcatV2", inputs=[
+                    "*",
+                    OpTypePattern("ExpandDims", inputs=[
+                        OpTypePattern("Minimum|Cast", name="min_or_cast"),
+                        "*"
+                    ]),
+                    "*",
+                ]),
+                OpTypePattern("Const", name="fill_value"),
+            ]), "*"
+        ])
 
-    for pattern in [pattern1, pattern2, pattern3, pattern4]:
+    for pattern in [pattern1, pattern2, pattern3, pattern4, pattern5, pattern6]:
         matcher = GraphMatcher(pattern, allow_reorder=True)
         match_results = list(matcher.match_ops(ops))
         for match_result in match_results: