Full conversion of MatrixBandPart

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -2979,6 +2979,35 @@ def func(input_x):
             return tf.identity(res, name=_TFOUTPUT), tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
+    @check_opset_min_version(11, "CumSum")
+    def test_matrix_band_part_3(self):
+        for low, high in [(-1, 3), (2, 3), (4, 3), (0, -1), (0, 0)]:
+            input_val = np.random.randint(0, 666, (10, 15)).astype(np.int32)
+            def func(input_x):
+                res = tf.linalg.band_part(input_x, low, high)
+                return tf.identity(res, name=_TFOUTPUT)
+            self._run_test_case(func, [_OUTPUT], {_INPUT: input_val})
+
+    @check_opset_min_version(11, "CumSum")
+    def test_matrix_band_part_4(self):
+        for low, high in [(-1, 3), (2, 3), (4, 3), (0, -1), (0, 0)]:
+            input_val = np.random.randint(0, 666, (2, 3, 10, 15)).astype(np.int32)
+            def func(input_x):
+                res = tf.linalg.band_part(input_x, low, high)
+                return tf.identity(res, name=_TFOUTPUT)
+            self._run_test_case(func, [_OUTPUT], {_INPUT: input_val})
+
+    @check_opset_min_version(11, "CumSum")
+    def test_matrix_band_part_5(self):
+        for low_val, high_val in [(2, 3), (4, 3), (0, 0), (2, 0)]:
+            low_val = np.array(low_val, np.int32)
+            high_val = np.array(high_val, np.int32)
+            input_val = np.random.randint(0, 666, (2, 3, 10, 15)).astype(np.int32)
+            def func(input_x, low, high):
+                res = tf.linalg.band_part(input_x, low, high)
+                return tf.identity(res, name=_TFOUTPUT)
+            self._run_test_case(func, [_OUTPUT], {_INPUT: input_val, _INPUT1: low_val, _INPUT2: high_val})
+
     def test_floordiv(self):
         input_val_1 = np.random.random_sample(100).astype(np.int32)
         input_val_2 = (np.random.random_sample(100) + 1).astype(np.int32)

tf2onnx/onnx_opset/nn.py

@@ -12,7 +12,7 @@
 import logging
 
 import numpy as np
-from onnx import onnx_pb
+from onnx import onnx_pb, helper
 from onnx.onnx_pb import TensorProto
 from tf2onnx import constants, utils
 from tf2onnx.graph_builder import GraphBuilder
@@ -1188,13 +1188,13 @@ def version_11(cls, ctx, node, **kwargs):
 @tf_op("MatrixBandPart")
 class MatrixBandPart:
     @classmethod
-    def any_version_after7(cls, opset, ctx, node, **kwargs):
+    def version_7(cls, opset, ctx, node, **kwargs):
         # T output = MatrixBandPart(T input, int num_lower, int num_upper)
         # data-flow: first generate mask matrix and then use element-wise mul op
         input_rank = len(ctx.get_shape(node.input[0]))
         utils.make_sure(input_rank == 2, error_msg="MatrixBandPart op: only rank 2 is supported")
         bandpart = [node.inputs[ind].get_tensor_value() for ind in [1, 2]]
-        utils.make_sure(bandpart in [[-1, 0], [0, -1]], "only support Lower/Upper triangular for now")
+        utils.make_sure(bandpart in [[-1, 0], [0, -1]], "only support Lower/Upper triangular for opset < 11")
         # methods to generate mask matrix: if lower triangular is needed, then generate column one by one
         # otherwise row is generated one by one.
         axis, counter_axis, squeeze_axis = (1, 0, 2) if bandpart == [-1, 0] else (0, 1, 1)
@@ -1267,13 +1267,76 @@ def any_version_after7(cls, opset, ctx, node, **kwargs):
                       dtypes=dtypes)
 
     @classmethod
-    def version_7(cls, ctx, node, **kwargs):
-        cls.any_version_after7(7, ctx, node, **kwargs)
-
-    @classmethod
-    def version_13(cls, ctx, node, **kwargs):
-        # Signature of operator Squeeze changed.
-        cls.any_version_after7(13, ctx, node, **kwargs)
+    def version_11(cls, ctx, node, **kwargs):
+        num_lower_const = node.inputs[1].get_tensor_value() if node.inputs[1].is_const() else None
+        num_upper_const = node.inputs[2].get_tensor_value() if node.inputs[2].is_const() else None
+        data, num_lower, num_upper = node.input
+        rank = ctx.get_rank(data)
+        int_max_val = utils.get_max_value(np.int64)
+        dtype = ctx.get_dtype(data)
+        np_dtype = utils.map_onnx_to_numpy_type(dtype)
+        if rank == 2:
+            shape = ctx.make_node("Shape", [data]).output[0]
+        else:
+            whole_shape = ctx.make_node("Shape", [data]).output[0]
+            shape = GraphBuilder(ctx).make_slice(
+                {'data': whole_shape, 'starts': [-2], 'ends': [int_max_val], 'axes': [0]})
+        if num_lower_const == 0 and num_upper_const == 0:
+            if rank == 2:
+                identity_node = ctx.make_node("EyeLike", [data]).output[0]
+            else:
+                zero_tensor = helper.make_tensor("value", dtype, dims=[1], vals=[0])
+                const_of_shape = ctx.make_node("ConstantOfShape", [shape], attr={'value': zero_tensor}).output[0]
+                identity_node = ctx.make_node("EyeLike", [const_of_shape]).output[0]
+            one_const = ctx.make_const(utils.make_name("one"), np.array(1, np_dtype)).output[0]
+            mask = ctx.make_node("Sub", [one_const, identity_node]).output[0]
+            shapes = node.output_shapes
+            dtypes = node.output_dtypes
+            ctx.remove_node(node.name)
+            ctx.make_node(op_type="Mul", inputs=[identity_node, data],
+                          name=node.name, outputs=node.output, shapes=shapes,
+                          dtypes=dtypes)
+            return
+        zero_const = ctx.make_const(utils.make_name("zero"), np.array(0, np.int64)).output[0]
+        one_const = ctx.make_const(utils.make_name("one"), np.array(1, np.int64)).output[0]
+        conditions = []
+        row_cnt = GraphBuilder(ctx).make_slice({'data': shape, 'axes': [0], 'starts': [0], 'ends': [1]})
+        col_cnt = GraphBuilder(ctx).make_slice({'data': shape, 'axes': [0], 'starts': [1], 'ends': [2]})
+        limit = ctx.make_node("Mul", [row_cnt, col_cnt]).output[0]
+        # idx_cnt = ctx.make_node("Range", [zero_const, limit, one_const]).output[0]
+
+        ones_of_shape = ctx.make_node("Expand", [one_const, limit]).output[0]
+        idx_cnt = ctx.make_node("CumSum", [ones_of_shape, zero_const], attr={'exclusive': True}).output[0]
+
+        idx_reshape = ctx.make_node("Reshape", [idx_cnt, shape]).output[0]
+        row_idx = ctx.make_node("Div", [idx_reshape, col_cnt]).output[0]
+        col_idx = ctx.make_node("Mod", [idx_reshape, col_cnt]).output[0]
+        idx_diff = ctx.make_node("Sub", [col_idx, row_idx]).output[0]
+
+        if num_upper_const is None or num_upper_const >= 0:
+            if ctx.get_dtype(num_upper) != TensorProto.INT64:
+                num_upper = ctx.make_node("Cast", [num_upper], attr={'to': TensorProto.INT64}).output[0]
+            conditions.append(ctx.make_node("LessOrEqual", [idx_diff, num_upper]).output[0])
+        if num_lower_const is None or num_lower_const >= 0:
+            if ctx.get_dtype(num_lower) != TensorProto.INT64:
+                num_lower = ctx.make_node("Cast", [num_lower], attr={'to': TensorProto.INT64}).output[0]
+            num_lower_neg = ctx.make_node("Neg", [num_lower]).output[0]
+            conditions.append(ctx.make_node("LessOrEqual", [num_lower_neg, idx_diff]).output[0])
+        if len(conditions) == 0:
+            node.type = "Identity"
+            ctx.replace_inputs(node, [data])
+            return
+        if len(conditions) == 1:
+            cond = conditions[0]
+        if len(conditions) == 2:
+            cond = ctx.make_node("And", conditions).output[0]
+        mask = ctx.make_node("Cast", [cond], attr={'to': ctx.get_dtype(data)}).output[0]
+        shapes = node.output_shapes
+        dtypes = node.output_dtypes
+        ctx.remove_node(node.name)
+        ctx.make_node(op_type="Mul", inputs=[mask, data],
+                      name=node.name, outputs=node.output, shapes=shapes,
+                      dtypes=dtypes)
 
 
 def _make_softmax_cross_entropy_with_logits(ctx, label, logit, tf_ori_node):