support MatrixDiagV1&V2

Author: RandySheriffH

tests/test_backend.py

@@ -3251,10 +3251,13 @@ def test_matrix_diag_v3_multi_dim(self):
 
         def func(diag, k, row, col):
             return tf.raw_ops.MatrixDiagV3(diagonal=diag, k=k, num_rows=row, num_cols=col,
-                                           padding_value=0.123, align='RIGHT_RIGHT', name=_TFOUTPUT)
+                                           padding_value=0.123, align='RIGHT_RIGHT', name=_TFOUTPUT), \
+                   tf.raw_ops.MatrixDiagV2(diagonal=diag, k=k, num_rows=row, num_cols=col,
+                                           padding_value=0.123, name=_TFOUTPUT1)
 
-        self._run_test_case(func, [_OUTPUT], {_INPUT: diag_val, _INPUT1: k_val,
-                                              _INPUT2: row_val, _INPUT3: col_val})
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1],
+                            {_INPUT: diag_val, _INPUT1: k_val,
+                            _INPUT2: row_val, _INPUT3: col_val})
 
     @check_opset_min_version(12)
     @check_tf_min_version("2.2")
@@ -3302,10 +3305,12 @@ def test_matrix_diag_v3_2single_dim_row_col(self):
 
         def func(diag, k, row, col):
             return tf.raw_ops.MatrixDiagV3(diagonal=diag, k=k, num_rows=row, num_cols=col,
-                                           padding_value=7, align='LEFT_RIGHT', name=_TFOUTPUT)
+                                           padding_value=7, align='LEFT_RIGHT', name=_TFOUTPUT), \
+                   tf.raw_ops.MatrixDiag(diagonal=diag, name=_TFOUTPUT1)
 
-        self._run_test_case(func, [_OUTPUT], {_INPUT: diag_val, _INPUT1: k_val,
-                                              _INPUT2: row_val, _INPUT3: col_val})
+        self._run_test_case(func, [_OUTPUT, _OUTPUT1],
+                            {_INPUT: diag_val, _INPUT1: k_val,
+                             _INPUT2: row_val, _INPUT3: col_val})
 
     @check_opset_min_version(12)
     @check_tf_min_version("2.2")

tf2onnx/onnx_opset/tensor.py

@@ -2207,11 +2207,13 @@ def compute_out_shape(k0_k1_same=False):
                 ctx.set_shape(consumer.output[0], shapes)
 
 
-@tf_op("MatrixDiagV3")
-class MatrixDiagV3:
+@tf_op(["MatrixDiag", "MatrixDiagV2", "MatrixDiagV3"])
+class MatrixDiag:
     @classmethod
     def version_12(cls, ctx, node, **kwargs):
         # Assemble MatrixDiagV3 by ReverseSequence
+        argc = len(node.input)
+
         def mkconsts(values):
             return [ctx.make_const(utils.make_name('const'), \
                                    np.array(value).astype(np.int64)).output[0] for value in values]
@@ -2230,6 +2232,9 @@ def normalize(name):
             reshaped = mknode("Reshape", [casted, minus_one])
             return reshaped
 
+        def cast(name):
+            return mknode("Cast", [name], attr={"to": ctx.get_dtype(node.input[0])})
+
         def processdiag():
             # unsqueeze diag if necessary
             diag = node.input[0]
@@ -2241,7 +2246,7 @@ def processdiag():
 
             diag_shape = mknode("Shape", [diag])
             diag_depth = mknode("Slice", [diag_shape, minus_two, minus_one])
-            k = normalize(node.input[1])
+            k = normalize(node.input[1]) if argc > 1 else zeo
             k_min, k_max = mknode("ReduceMin", [k]), mknode("ReduceMax", [k])
             k_max_nxt = mknode("Add", [k_max, one])
             k_depth = mknode("Sub", [k_max_nxt, k_min])
@@ -2272,8 +2277,10 @@ def squeeze(name):
 
         # gather inputs
         diag, k, k_min, k_max, k_max_nxt = processdiag()
-        row, col, pad, align = normalize(node.input[2]), normalize(node.input[3]), \
-                               node.input[4], node.get_attr_str("align")
+        row, col, pad, align = normalize(node.input[2]) if argc > 2 else minus_one, \
+                               normalize(node.input[3]) if argc > 3 else minus_one, \
+                               node.input[4] if argc > 4 else cast(zeo), \
+                               node.get_attr_str("align") if "align" in node.attr else "LEFT_LEFT"
 
         diag_shape = mknode("Shape", [diag])
         diag_rank = mknode("Shape", [diag_shape])
@@ -2580,12 +2587,12 @@ def normalize():
         # make matrix of bool
         ctx.set_dtype(ones_diag.output[0], TensorProto.INT64)
         ones_matrix = ctx.make_node("MatrixDiagV3", [ones_diag.output[0], k, row, col, zeo], attr)
-        MatrixDiagV3.version_12(ctx, ones_matrix)
+        MatrixDiag.version_12(ctx, ones_matrix)
         ones_bool = mknode("Equal", [ones_matrix.output[0], one])
 
         # make matrix out of diag
         diag_matrix = ctx.make_node("MatrixDiagV3", [diag, k, row, col, cast(zeo)], attr)
-        MatrixDiagV3.version_12(ctx, diag_matrix)
+        MatrixDiag.version_12(ctx, diag_matrix)
 
         shapes = node.output_shapes
         dtypes = node.output_dtypes