Merge pull request #1233 from onnx/tom/opset13updates

Update split op for opset13 and add test cases

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1366,6 +1366,15 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(13, "Split")
+    def test_split_nonconst(self):
+        x_val = np.linspace(1.0, 5 * 30.0, 5 * 30).astype(np.float32).reshape((5, 30))
+        y_val = np.array([4, 15, 11], np.int32)
+        def func(x, y):
+            x_, _, _ = tf.split(x, y, 1)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
     def test_split_with_more_outputs(self):
         x_val = np.linspace(1.0, 5 * 30.0, 5 * 30).astype(np.float32).reshape((5, 30))
         def func(x):
@@ -1387,6 +1396,24 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(13, "ReduceSum")
+    def test_reducesum_nonconst_axis(self):
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 1, 2))
+        y_val = np.array([1, 2], dtype=np.int32)
+        def func(x, y):
+            x_ = tf.reduce_sum(x, axis=y)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
+    @check_opset_min_version(13, "ReduceSum")
+    def test_reducesum_empty_axis(self):
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 1, 2))
+        y_val = np.array([], dtype=np.int32)
+        def func(x, y):
+            x_ = tf.reduce_sum(x, axis=y)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
     @check_opset_min_version(9, "OneHot")
     def test_segment_sum_data_vector(self):
         segs_val = np.array([0, 0, 0, 1, 2, 2, 3, 3], dtype=np.int32)
@@ -2866,6 +2893,16 @@ def func(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, "ReduceSum")
+    @check_tf_min_version("1.15")
+    def test_reduce_any_empty_axis(self):
+        input_val = np.random.randint(0, 2, (10, 20)).astype(np.bool)
+        def func(x):
+            res = tf.reduce_any(input_tensor=x, keepdims=False)
+            res1 = tf.reduce_any(input_tensor=x, axis=[], keepdims=False)
+            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(7, "fill")
     def test_zeros_like(self):
         input_x = np.random.random_sample([10, 20]).astype(np.float32)
@@ -3289,6 +3326,20 @@ def func(x):
             return tf.identity(y, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    def test_softmax(self):
+        x_val = np.arange(0, 24, dtype=np.float32).reshape([3, 1, 8])
+        def func(x):
+            y = tf.nn.softmax(x)
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    def test_log_softmax(self):
+        x_val = np.arange(0, 24, dtype=np.float32).reshape([3, 1, 8])
+        def func(x):
+            y = tf.nn.log_softmax(x)
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     # test for gemm pattern0: alpha*A*B + beta*C
     def test_gemm_pattern0(self):
         max_number = 10

tf2onnx/onnx_opset/math.py

@@ -201,6 +201,11 @@ def version_1(cls, ctx, node, **kwargs):
     def version_11(cls, ctx, node, **kwargs):
         cls.version_1(ctx, node, **kwargs)
 
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        # Default axis is now -1.
+        pass
+
 
 @tf_op("Square")
 class Square:

tf2onnx/onnx_opset/rnn.py

@@ -94,10 +94,9 @@ def make_sigmoid(i, w, b):
             raise RuntimeError("shape of W of LSTMBlockCell {} should be times of 4".format(node.name))
         merged_output_node = ctx.make_node("Add", [xh_w_node.output[0], b])
         w_last_dim = int(w_shape[1] / 4)
-        split = [w_last_dim] * 4
         split_output_node = ctx.make_node(
             "Split", [merged_output_node.output[0]],
-            attr={"axis": 1, "split": split},
+            attr={"axis": 1},
             output_count=4
         )
         i, ci, f, o = split_output_node.output

tf2onnx/onnx_opset/tensor.py

@@ -625,6 +625,10 @@ def version_11(cls, ctx, node, **kwargs):
         # no change
         cls.version_1(ctx, node, **kwargs)
 
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        # Default axis is not -1 but doesn't matter since we always set it.
+        cls.version_1(ctx, node, **kwargs)
 
 @tf_op("SplitV")
 class SplitV:
@@ -652,6 +656,26 @@ def version_1(cls, ctx, node, **kwargs):
     def version_2(cls, ctx, node, **kwargs):
         cls.version_1(ctx, node, **kwargs)
 
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        # Split now supports dynamic split lengths
+        if node.inputs[1].is_const():
+            # Call version 1 to deal with -1 cases
+            cls.version_1(ctx, node, **kwargs)
+            # Convert attr to input
+            split_val = node.get_attr_value("split")
+            split_const = ctx.make_const(utils.make_name("split"), np.array(split_val, np.int64))
+            ctx.replace_inputs(node, [node.input[0], split_const.output[0]])
+            del node.attr["split"]
+        else:
+            # Technically incorrect if any of the splits are -1
+            node.type = "Split"
+            split_dims = node.inputs[2].get_tensor_value()
+            ctx.remove_input(node, node.input[2], 2)
+            node.set_attr("axis", split_dims)
+            if ctx.get_dtype(node.input[1]) != TensorProto.INT64:
+                ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=TensorProto.INT64)
+
 
 @tf_op("ExpandDims")
 class ExpandDims: