Enhance BatchToSpaceND to support 3D input data

tests/test_backend.py

@@ -2359,6 +2359,15 @@ def test_batch_to_spacend(self):
         _ = tf.batch_to_space_nd(input_x, block_size, crop, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: input_val})
 
+    def test_batch_to_space3d(self):
+        block_size = [2, 2]
+        crop = [[0, 1], [2, 1]]
+
+        input_val = np.random.random_sample([40, 3, 100]).astype(np.float32)
+        input_x = tf.placeholder(dtype=tf.float32, shape=input_val.shape, name=_TFINPUT)  # NHC
+        _ = tf.batch_to_space_nd(input_x, block_size, crop, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: input_val})
+
     def test_space_to_batchnd(self):
         block_size = [2, 2]
         pad = [[0, 1], [2, 1]]

tf2onnx/onnx_opset/tensor.py

@@ -1068,20 +1068,26 @@ class BatchToSpace:
     def version_1(cls, ctx, node, **kwargs):
         # https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/batch-to-space-n-d.html
         # the above link says the data format of input tensor should be (batch, spatial_shape, remaining_shape)
-        # and we only support 4D here, so the data format is NHWC
+        # and we only support 3D and 4D here, and the data format is NHC and NHWC
         # onnx op "DepthToSpace" does the same work on input tensor except that it works on "C",
         # and it only supports NCHW
         # T out = BatchToSpaceND(T input, int32 block_shape, int32 crops)
         input_tensor = node.inputs[0]
         blocksize = node.inputs[1].get_tensor_value()
         crops = node.inputs[2].get_tensor_value()
 
-        utils.make_sure(len(ctx.get_shape(input_tensor.output[0])) == 4, "only supports 4D for now")
+        utils.make_sure(len(ctx.get_shape(input_tensor.output[0])) in (4, 3),
+                        "only supports 3D and 4D for now")
         utils.make_sure(len(blocksize) == 2 and blocksize[0] == blocksize[1],
                         "only support same blocksize at different dims")
 
         # NHWC TO CNHW, so onnx op will work on "N" which is the same as tensorflow
-        trans1 = ctx.make_node("Transpose", input_tensor.output, {"perm": [3, 0, 1, 2]})
+        if len(ctx.get_shape(input_tensor.output[0])) == 3:
+            # insert automatically an Unsqueeze op if the input is 3d
+            unsqz1 = ctx.make_node("Unsqueeze", input_tensor.output, {"axes": [3]})
+            trans1 = ctx.make_node("Transpose", unsqz1.output, {"perm": [3, 0, 1, 2]})
+        else:
+            trans1 = ctx.make_node("Transpose", input_tensor.output, {"perm": [3, 0, 1, 2]})
         reorganize_node = ctx.make_node(node.type, trans1.output, attr={"blocksize": blocksize[0]})
         trans2 = ctx.make_node("Transpose", reorganize_node.output, {"perm": [1, 2, 3, 0]})
 
@@ -1099,11 +1105,19 @@ def version_1(cls, ctx, node, **kwargs):
 
         attr = {"axes": slice_axis, "ends": ends, "starts": starts}
         inputs_map = {"data": trans2.output[0], **attr}
-        kwargs = {**inputs_map, "outputs": node.output}
         dtypes = [ctx.get_dtype(node.output[0])]
-        shapes = [ctx.get_shape(node.output[0])]
-        ctx.remove_node(node.name)
-        GraphBuilder(ctx).make_slice(kwargs, name=node.name, dtypes=dtypes, shapes=shapes)
+        shapes = ctx.get_shape(node.output[0])
+
+        if len(ctx.get_shape(input_tensor.output[0])) == 3:
+            # add a squeeze op to convert output into 3d
+            kwargs = {**inputs_map}
+            ctx.remove_node(node.name)
+            slice1 = GraphBuilder(ctx).make_slice(kwargs)
+            ctx.make_node("Squeeze", [slice1], {"axes": [3]}, outputs=node.output, name=node.name, dtypes=dtypes)
+        else:
+            kwargs = {**inputs_map, "outputs": node.output}
+            ctx.remove_node(node.name)
+            GraphBuilder(ctx).make_slice(kwargs, name=node.name, dtypes=dtypes, shapes=[shapes])
 
 
 @tf_op("SpaceToBatchND", onnx_op="SpaceToDepth")