basic support for tf.stack

Author: guschmue

tests/test_backend.py

@@ -834,6 +834,22 @@ def test_topk(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    def test_stack_axis0(self):
+        x_val = [np.random.randn(3, 4).astype("float32") for _ in range(10)]
+        x = [tf.constant(x_val[i], dtype=tf.float32) for i in range(10)]
+        x_ = tf.stack(x, axis=0)
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {}, {})
+        self.assertAllClose(expected, actual)
+
+    def test_stack_axis1(self):
+        x_val = [np.random.randn(3, 4).astype("float32") for _ in range(10)]
+        x = [tf.constant(x_val[i], dtype=tf.float32) for i in range(10)]
+        x_ = tf.stack(x, axis=1)
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {}, {})
+        self.assertAllClose(expected, actual)
+
     @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "Space2Depth not implemented, works on onnxmsrtnext")
     def test_space_to_depth(self):
         x_val = make_xval([1, 2, 2, 1])
@@ -873,7 +889,6 @@ def test_strided_slice2(self):
 
     @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "not correctly supported")
     def test_resize_nearest_neighbor(self):
-        # this should work but no runtime I tried supports it.
         x_shape = [1, 15, 20, 2]
         x_new_size = [30, 40]
         x_val = np.arange(1, 1 + np.prod(x_shape)).astype("float32").reshape(x_shape)
@@ -886,7 +901,6 @@ def test_resize_nearest_neighbor(self):
 
     @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "not correctly supported")
     def test_resize_bilinear(self):
-        # this should work but no runtime I tried supports it.
         x_shape = [1, 15, 20, 2]
         x_new_size = [30, 40]
         x_val = np.arange(1, 1 + np.prod(x_shape)).astype("float32").reshape(x_shape)

tf2onnx/tfonnx.py

@@ -832,6 +832,27 @@ def minmax_op(ctx, node, name, args):
     return node
 
 
+def pack_op(ctx, node, name, args):
+    # hack to make up for the missing onnx pack op
+    axis = node.get_attr("axis").i
+    nodes = []
+    inputs = []
+    # insert Unsqueeze on each input
+    for i, n in enumerate(node.inputs):
+        op_name = utils.make_name(node.name)
+        output_name = op_name + ":0"
+        new_node = Node(helper.make_node("Unsqueeze", [node.input[i]], [output_name], name=op_name, axes=[axis]), ctx)
+        node.input[i] = output_name
+        nodes.append(new_node)
+        inputs.append(output_name)
+    # concat all unqueezes
+    op_name = utils.make_name(node.name)
+    output_name = op_name + ":0"
+    concat = Node(helper.make_node("Concat", inputs, [output_name], name=op_name, axis=axis), ctx)
+    ctx.copy_shape(node.output[0], concat.output[0])
+    ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], output_name)
+    return [concat] + nodes
+
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -917,6 +938,7 @@ def minmax_op(ctx, node, name, args):
     "Transpose": (transpose_op, []),
     "TopKV2": (topk_op, []),
     "SpaceToDepth": (spacetodepth_op, []),
+    "Pack": (pack_op, []),
 }
 
 _OPSET_5 = {