Merge pull request #891 from onnx/gs/fix-negative-split

fix split in case of splits are negavitve

Author: guschmue

tests/test_backend.py

@@ -1160,6 +1160,13 @@ def func(x):
             return tf.split(x, [4, 15, 11], 1, name="split_test")
         self._run_test_case(func, ["split_test:0", "split_test:1", "split_test:2"], {_INPUT: x_val})
 
+    def test_negative_split(self):
+        x_val = np.linspace(1.0, 5 * 30.0, 5 * 30).astype(np.float32).reshape((5, 30))
+        def func(x):
+            x_, _, _ = tf.split(x, [4, 15, -1], 1)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     def test_reducesum(self):
         # not supported by onnx-caffe2
         x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2))

tf2onnx/onnx_opset/tensor.py

@@ -553,6 +553,14 @@ def version_1(cls, ctx, node, **kwargs):
         node.type = "Split"
         split = node.inputs[1].get_tensor_value()
         split_dims = node.inputs[2].get_tensor_value()
+        if -1 in split:
+            # negative split = use the remaining size
+            shape = ctx.get_shape(node.input[0])
+            final_sum = shape[split_dims]
+            sums = sum([i for i in split if i >= 0])
+            for i, v in enumerate(split):
+                if v == -1:
+                    split[i] = final_sum - sums
         ctx.remove_input(node, node.input[2])
         ctx.remove_input(node, node.input[1])
         node.set_attr("split", split)